From f26e21c1a13a062f26607cff201dccb3359f5250 Mon Sep 17 00:00:00 2001
From: Minh Vu <vuhoangminh97@gmail.com>
Date: Wed, 15 Jul 2026 05:38:34 +0200
Subject: [PATCH] fix(flight): reject malformed basic credentials

---
 arrow/flight/basic_auth_flight_test.go | 30 +++++++++++++++++++++++++-
 arrow/flight/server_auth.go            |  7 ++++--
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/arrow/flight/basic_auth_flight_test.go b/arrow/flight/basic_auth_flight_test.go
index 744ea29f..72c64575 100644
--- a/arrow/flight/basic_auth_flight_test.go
+++ b/arrow/flight/basic_auth_flight_test.go
@@ -18,6 +18,7 @@ package flight_test
 
 import (
 	"context"
+	"encoding/base64"
 	"errors"
 	"io"
 	"testing"
@@ -32,7 +33,7 @@ import (
 
 const (
 	validUsername   = "flight_username"
-	validPassword   = "flight_password"
+	validPassword   = "flight_password:with_colon"
 	invalidUsername = "invalid_flight_username"
 	invalidPassword = "invalid_flight_password"
 	validBearer     = "CAREBARESTARE"
@@ -249,3 +250,30 @@ func TestBasicAuthMissingCredential(t *testing.T) {
 		t.Fatalf("unexpected code: got %v, want %v", got, want)
 	}
 }
+
+func TestBasicAuthMalformedCredentials(t *testing.T) {
+	s := flight.NewServerWithMiddleware([]flight.ServerMiddleware{flight.CreateServerBasicAuthMiddleware(&validator{})})
+	s.Init("localhost:0")
+	s.RegisterFlightService(&HeaderAuthTestFlight{})
+	go s.Serve()
+	defer s.Shutdown()
+
+	client, err := flight.NewFlightClient(s.Addr().String(), nil, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	for _, credentials := range []string{"", "username-only", ":password", "username:"} {
+		t.Run(credentials, func(t *testing.T) {
+			encoded := base64.StdEncoding.EncodeToString([]byte(credentials))
+			ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs("authorization", "Basic "+encoded))
+			stream, err := client.Handshake(ctx)
+			if err == nil {
+				_, err = stream.Recv()
+			}
+			if status.Code(err) != codes.Unauthenticated {
+				t.Fatalf("expected unauthenticated error, got %v", err)
+			}
+		})
+	}
+}
diff --git a/arrow/flight/server_auth.go b/arrow/flight/server_auth.go
index 7135e43a..155c2ac6 100644
--- a/arrow/flight/server_auth.go
+++ b/arrow/flight/server_auth.go
@@ -195,8 +195,11 @@ func createServerBearerTokenStreamInterceptor(validator BasicAuthValidator) grpc
 					}
 				}
 
-				creds := strings.SplitN(string(val), ":", 2)
-				token, err := validator.Validate(creds[0], creds[1])
+				username, password, found := strings.Cut(string(val), ":")
+				if !found {
+					return status.Error(codes.Unauthenticated, "invalid basic auth credentials")
+				}
+				token, err := validator.Validate(username, password)
 				if err != nil {
 					return err
 				}