Skip to content

Commit 1ed6b9e

Browse files
mudlermudler
committed
fix(llama.cpp): correctly parse grpc header for bearer token auth
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>
1 parent e4ee743 commit 1ed6b9e

1 file changed

Lines changed: 48 additions & 45 deletions

File tree

‎backend/cpp/llama-cpp/grpc-server.cpp‎

Lines changed: 48 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -40,45 +40,41 @@ using grpc::ServerBuilder;
4040
using grpc::ServerContext;
4141
using grpc::Status;
4242

43-
// gRPC bearer token auth via AuthMetadataProcessor for distributed mode.
43+
// gRPC bearer token auth for distributed mode.
4444
// Reads LOCALAI_GRPC_AUTH_TOKEN from the environment. When set, rejects
4545
// requests without a matching "authorization: Bearer <token>" metadata header.
46-
class TokenAuthMetadataProcessor : public grpc::AuthMetadataProcessor {
47-
public:
48-
explicit TokenAuthMetadataProcessor(const std::string& token) : token_(token) {}
49-
50-
bool IsBlocking() const override { return false; }
51-
52-
grpc::Status Process(const InputMetadata& auth_metadata,
53-
grpc::AuthContext* /*context*/,
54-
OutputMetadata* /*consumed_auth_metadata*/,
55-
OutputMetadata* /*response_metadata*/) override {
56-
auto it = auth_metadata.find("authorization");
57-
if (it != auth_metadata.end()) {
58-
std::string expected = "Bearer " + token_;
59-
std::string got(it->second.data(), it->second.size());
60-
// Constant-time comparison
61-
if (expected.size() == got.size() && ct_memcmp(expected.data(), got.data(), expected.size()) == 0) {
62-
return grpc::Status::OK;
63-
}
64-
}
65-
return grpc::Status(grpc::StatusCode::UNAUTHENTICATED, "invalid token");
46+
47+
// Cached auth token — empty means auth is disabled.
48+
static std::string g_grpc_auth_token;
49+
50+
// Minimal constant-time comparison (avoids OpenSSL dependency)
51+
static int ct_memcmp(const void* a, const void* b, size_t n) {
52+
const unsigned char* pa = static_cast<const unsigned char*>(a);
53+
const unsigned char* pb = static_cast<const unsigned char*>(b);
54+
unsigned char result = 0;
55+
for (size_t i = 0; i < n; i++) {
56+
result |= pa[i] ^ pb[i];
6657
}
58+
return result;
59+
}
6760

68-
private:
69-
std::string token_;
70-
71-
// Minimal constant-time comparison (avoids OpenSSL dependency)
72-
static int ct_memcmp(const void* a, const void* b, size_t n) {
73-
const unsigned char* pa = static_cast<const unsigned char*>(a);
74-
const unsigned char* pb = static_cast<const unsigned char*>(b);
75-
unsigned char result = 0;
76-
for (size_t i = 0; i < n; i++) {
77-
result |= pa[i] ^ pb[i];
61+
// Returns OK when auth is disabled or the token matches.
62+
static grpc::Status checkAuth(grpc::ServerContext* context) {
63+
if (g_grpc_auth_token.empty()) {
64+
return grpc::Status::OK;
65+
}
66+
auto metadata = context->client_metadata();
67+
auto it = metadata.find("authorization");
68+
if (it != metadata.end()) {
69+
std::string expected = "Bearer " + g_grpc_auth_token;
70+
std::string got(it->second.data(), it->second.size());
71+
if (expected.size() == got.size() &&
72+
ct_memcmp(expected.data(), got.data(), expected.size()) == 0) {
73+
return grpc::Status::OK;
7874
}
79-
return result;
8075
}
81-
};
76+
return grpc::Status(grpc::StatusCode::UNAUTHENTICATED, "invalid token");
77+
}
8278

8379
// END LocalAI
8480

@@ -757,13 +753,17 @@ class BackendServiceImpl final : public backend::Backend::Service {
757753
public:
758754
BackendServiceImpl(server_context& ctx) : ctx_server(ctx) {}
759755

760-
grpc::Status Health(ServerContext* /*context*/, const backend::HealthMessage* /*request*/, backend::Reply* reply) override {
756+
grpc::Status Health(ServerContext* context, const backend::HealthMessage* /*request*/, backend::Reply* reply) override {
757+
auto auth = checkAuth(context);
758+
if (!auth.ok()) return auth;
761759
// Implement Health RPC
762760
reply->set_message("OK");
763761
return Status::OK;
764762
}
765763

766-
grpc::Status LoadModel(ServerContext* /*context*/, const backend::ModelOptions* request, backend::Result* result) override {
764+
grpc::Status LoadModel(ServerContext* context, const backend::ModelOptions* request, backend::Result* result) override {
765+
auto auth = checkAuth(context);
766+
if (!auth.ok()) return auth;
767767
// Implement LoadModel RPC
768768
common_params params;
769769
params_parse(ctx_server, request, params);
@@ -962,6 +962,8 @@ class BackendServiceImpl final : public backend::Backend::Service {
962962
}
963963

964964
grpc::Status PredictStream(grpc::ServerContext* context, const backend::PredictOptions* request, grpc::ServerWriter<backend::Reply>* writer) override {
965+
auto auth = checkAuth(context);
966+
if (!auth.ok()) return auth;
965967
if (params_base.model.path.empty()) {
966968
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
967969
}
@@ -1665,6 +1667,8 @@ class BackendServiceImpl final : public backend::Backend::Service {
16651667
}
16661668

16671669
grpc::Status Predict(ServerContext* context, const backend::PredictOptions* request, backend::Reply* reply) override {
1670+
auto auth = checkAuth(context);
1671+
if (!auth.ok()) return auth;
16681672
if (params_base.model.path.empty()) {
16691673
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
16701674
}
@@ -2383,6 +2387,8 @@ class BackendServiceImpl final : public backend::Backend::Service {
23832387
}
23842388

23852389
grpc::Status Embedding(ServerContext* context, const backend::PredictOptions* request, backend::EmbeddingResult* embeddingResult) override {
2390+
auto auth = checkAuth(context);
2391+
if (!auth.ok()) return auth;
23862392
if (params_base.model.path.empty()) {
23872393
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
23882394
}
@@ -2563,7 +2569,9 @@ class BackendServiceImpl final : public backend::Backend::Service {
25632569
return grpc::Status::OK;
25642570
}
25652571

2566-
grpc::Status TokenizeString(ServerContext* /*context*/, const backend::PredictOptions* request, backend::TokenizationResponse* response) override {
2572+
grpc::Status TokenizeString(ServerContext* context, const backend::PredictOptions* request, backend::TokenizationResponse* response) override {
2573+
auto auth = checkAuth(context);
2574+
if (!auth.ok()) return auth;
25672575
if (params_base.model.path.empty()) {
25682576
return grpc::Status(grpc::StatusCode::FAILED_PRECONDITION, "Model not loaded");
25692577
}
@@ -2803,19 +2811,14 @@ int main(int argc, char** argv) {
28032811
BackendServiceImpl service(ctx_server);
28042812

28052813
ServerBuilder builder;
2806-
// Add bearer token auth via AuthMetadataProcessor if LOCALAI_GRPC_AUTH_TOKEN is set
2814+
builder.AddListeningPort(server_address, grpc::InsecureServerCredentials());
2815+
2816+
// Initialize bearer token auth if LOCALAI_GRPC_AUTH_TOKEN is set
28072817
const char* auth_token = std::getenv("LOCALAI_GRPC_AUTH_TOKEN");
2808-
std::shared_ptr<grpc::ServerCredentials> creds;
28092818
if (auth_token != nullptr && auth_token[0] != '\0') {
2810-
creds = grpc::InsecureServerCredentials();
2811-
creds->SetAuthMetadataProcessor(
2812-
std::make_shared<TokenAuthMetadataProcessor>(auth_token));
2819+
g_grpc_auth_token = auth_token;
28132820
std::cout << "gRPC auth enabled via LOCALAI_GRPC_AUTH_TOKEN" << std::endl;
2814-
} else {
2815-
creds = grpc::InsecureServerCredentials();
28162821
}
2817-
2818-
builder.AddListeningPort(server_address, creds);
28192822
builder.RegisterService(&service);
28202823
builder.SetMaxMessageSize(50 * 1024 * 1024); // 50MB
28212824
builder.SetMaxSendMessageSize(50 * 1024 * 1024); // 50MB

0 commit comments

Comments
 (0)