2222#include < grpcpp/ext/proto_server_reflection_plugin.h>
2323#include < grpcpp/grpcpp.h>
2424#include < grpcpp/health_check_service_interface.h>
25+ #include < grpcpp/security/server_credentials.h>
2526#include < regex>
2627#include < atomic>
28+ #include < cstdlib>
2729#include < mutex>
2830#include < signal.h>
2931#include < thread>
@@ -37,6 +39,47 @@ using grpc::Server;
3739using grpc::ServerBuilder;
3840using grpc::ServerContext;
3941using grpc::Status;
42+
43+ // gRPC bearer token auth via AuthMetadataProcessor for distributed mode.
44+ // Reads LOCALAI_GRPC_AUTH_TOKEN from the environment. When set, rejects
45+ // 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" );
66+ }
67+
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];
78+ }
79+ return result;
80+ }
81+ };
82+
4083// END LocalAI
4184
4285
@@ -2760,11 +2803,24 @@ int main(int argc, char** argv) {
27602803 BackendServiceImpl service (ctx_server);
27612804
27622805 ServerBuilder builder;
2763- builder.AddListeningPort (server_address, grpc::InsecureServerCredentials ());
2806+ // Add bearer token auth via AuthMetadataProcessor if LOCALAI_GRPC_AUTH_TOKEN is set
2807+ const char * auth_token = std::getenv (" LOCALAI_GRPC_AUTH_TOKEN" );
2808+ std::shared_ptr<grpc::ServerCredentials> creds;
2809+ if (auth_token != nullptr && auth_token[0 ] != ' \0 ' ) {
2810+ creds = grpc::InsecureServerCredentials ();
2811+ creds->SetAuthMetadataProcessor (
2812+ std::make_shared<TokenAuthMetadataProcessor>(auth_token));
2813+ std::cout << " gRPC auth enabled via LOCALAI_GRPC_AUTH_TOKEN" << std::endl;
2814+ } else {
2815+ creds = grpc::InsecureServerCredentials ();
2816+ }
2817+
2818+ builder.AddListeningPort (server_address, creds);
27642819 builder.RegisterService (&service);
27652820 builder.SetMaxMessageSize (50 * 1024 * 1024 ); // 50MB
27662821 builder.SetMaxSendMessageSize (50 * 1024 * 1024 ); // 50MB
27672822 builder.SetMaxReceiveMessageSize (50 * 1024 * 1024 ); // 50MB
2823+
27682824 std::unique_ptr<Server> server (builder.BuildAndStart ());
27692825 // run the HTTP server in a thread - see comment below
27702826 std::thread t ([&]()
0 commit comments