Skip to content

Commit 359fbe3

Browse files
acsotoacsoto
committed
fix(router): add auth header for workload manager
Signed-off-by: Zhou Zihang <z@mcac.cc>
1 parent eb947f5 commit 359fbe3

File tree

2 files changed

+184
-0
lines changed

2 files changed

+184
-0
lines changed

‎pkg/router/session_manager.go‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,8 +32,12 @@ import (
3232
"github.com/volcano-sh/agentcube/pkg/common/types"
3333
"github.com/volcano-sh/agentcube/pkg/store"
3434
"golang.org/x/net/http2"
35+
"k8s.io/klog/v2"
3536
)
3637

38+
// #nosec G101 -- well-known Kubernetes service account token path, not a credential.
39+
var serviceAccountTokenPath = "/var/run/secrets/kubernetes.io/serviceaccount/token"
40+
3741
// SessionManager defines the session management behavior on top of Store and the workload manager.
3842
type SessionManager interface {
3943
// GetSandboxBySession returns the sandbox associated with the given sessionID.
@@ -139,6 +143,9 @@ func (m *manager) createSandbox(ctx context.Context, namespace string, name stri
139143
return nil, fmt.Errorf("failed to create HTTP request: %w", err)
140144
}
141145
req.Header.Set("Content-Type", "application/json")
146+
if token := loadWorkloadManagerAuthToken(); token != "" {
147+
req.Header.Set("Authorization", "Bearer "+token)
148+
}
142149

143150
// Send the request
144151
resp, err := m.httpClient.Do(req)
@@ -186,3 +193,18 @@ func (m *manager) createSandbox(ctx context.Context, namespace string, name stri
186193

187194
return sandbox, nil
188195
}
196+
197+
func loadWorkloadManagerAuthToken() string {
198+
if token := strings.TrimSpace(os.Getenv("API_TOKEN")); token != "" {
199+
return token
200+
}
201+
202+
b, err := os.ReadFile(serviceAccountTokenPath)
203+
if err != nil {
204+
if !os.IsNotExist(err) {
205+
klog.Warningf("failed to read service account token: %v", err)
206+
}
207+
return ""
208+
}
209+
return strings.TrimSpace(string(b))
210+
}

‎pkg/router/session_manager_test.go‎

Lines changed: 162 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,8 @@ import (
2222
"io"
2323
"net/http"
2424
"net/http/httptest"
25+
"os"
26+
"path/filepath"
2527
"testing"
2628
"time"
2729

@@ -226,6 +228,166 @@ func TestGetSandboxBySession_CreateSandbox_AgentRuntime_Success(t *testing.T) {
226228
}
227229
}
228230

231+
func TestGetSandboxBySession_CreateSandbox_SetsAuthHeaderFromEnv(t *testing.T) {
232+
// #nosec G101 -- test token, not a real credential.
233+
const token = "env-token-123"
234+
235+
if err := os.Setenv("API_TOKEN", token); err != nil {
236+
t.Fatalf("failed to set API_TOKEN: %v", err)
237+
}
238+
defer func() {
239+
_ = os.Unsetenv("API_TOKEN")
240+
}()
241+
242+
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
243+
if got := r.Header.Get("Authorization"); got != "Bearer "+token {
244+
t.Fatalf("expected Authorization header %q, got %q", "Bearer "+token, got)
245+
}
246+
resp := types.CreateSandboxResponse{
247+
SessionID: "new-session-123",
248+
SandboxID: "sandbox-456",
249+
SandboxName: "sandbox-test",
250+
EntryPoints: []types.SandboxEntryPoint{{Endpoint: "10.0.0.1:9000"}},
251+
}
252+
w.Header().Set("Content-Type", "application/json")
253+
w.WriteHeader(http.StatusOK)
254+
_ = json.NewEncoder(w).Encode(resp)
255+
}))
256+
defer mockServer.Close()
257+
258+
m := &manager{
259+
storeClient: &fakeStoreClient{},
260+
workloadMgrAddr: mockServer.URL,
261+
httpClient: &http.Client{},
262+
}
263+
264+
if _, err := m.GetSandboxBySession(context.Background(), "", "default", "test-runtime", types.AgentRuntimeKind); err != nil {
265+
t.Fatalf("GetSandboxBySession unexpected error: %v", err)
266+
}
267+
}
268+
269+
func TestGetSandboxBySession_CreateSandbox_SetsAuthHeaderFromFile(t *testing.T) {
270+
// #nosec G101 -- test token, not a real credential.
271+
const token = "file-token-456"
272+
273+
dir := t.TempDir()
274+
tokenPath := filepath.Join(dir, "token")
275+
if err := os.WriteFile(tokenPath, []byte(token+"\n"), 0o600); err != nil {
276+
t.Fatalf("failed to write token file: %v", err)
277+
}
278+
279+
origTokenPath := serviceAccountTokenPath
280+
serviceAccountTokenPath = tokenPath
281+
defer func() {
282+
serviceAccountTokenPath = origTokenPath
283+
}()
284+
285+
_ = os.Unsetenv("API_TOKEN")
286+
287+
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
288+
if got := r.Header.Get("Authorization"); got != "Bearer "+token {
289+
t.Fatalf("expected Authorization header %q, got %q", "Bearer "+token, got)
290+
}
291+
resp := types.CreateSandboxResponse{
292+
SessionID: "new-session-123",
293+
SandboxID: "sandbox-456",
294+
SandboxName: "sandbox-test",
295+
EntryPoints: []types.SandboxEntryPoint{{Endpoint: "10.0.0.1:9000"}},
296+
}
297+
w.Header().Set("Content-Type", "application/json")
298+
w.WriteHeader(http.StatusOK)
299+
_ = json.NewEncoder(w).Encode(resp)
300+
}))
301+
defer mockServer.Close()
302+
303+
m := &manager{
304+
storeClient: &fakeStoreClient{},
305+
workloadMgrAddr: mockServer.URL,
306+
httpClient: &http.Client{},
307+
}
308+
309+
if _, err := m.GetSandboxBySession(context.Background(), "", "default", "test-runtime", types.AgentRuntimeKind); err != nil {
310+
t.Fatalf("GetSandboxBySession unexpected error: %v", err)
311+
}
312+
}
313+
314+
func TestGetSandboxBySession_CreateSandbox_NoAuthHeaderWhenNoToken(t *testing.T) {
315+
dir := t.TempDir()
316+
missingPath := filepath.Join(dir, "missing-token")
317+
318+
origTokenPath := serviceAccountTokenPath
319+
serviceAccountTokenPath = missingPath
320+
defer func() {
321+
serviceAccountTokenPath = origTokenPath
322+
}()
323+
324+
_ = os.Unsetenv("API_TOKEN")
325+
326+
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
327+
if got := r.Header.Get("Authorization"); got != "" {
328+
t.Fatalf("expected no Authorization header, got %q", got)
329+
}
330+
resp := types.CreateSandboxResponse{
331+
SessionID: "new-session-123",
332+
SandboxID: "sandbox-456",
333+
SandboxName: "sandbox-test",
334+
EntryPoints: []types.SandboxEntryPoint{{Endpoint: "10.0.0.1:9000"}},
335+
}
336+
w.Header().Set("Content-Type", "application/json")
337+
w.WriteHeader(http.StatusOK)
338+
_ = json.NewEncoder(w).Encode(resp)
339+
}))
340+
defer mockServer.Close()
341+
342+
m := &manager{
343+
storeClient: &fakeStoreClient{},
344+
workloadMgrAddr: mockServer.URL,
345+
httpClient: &http.Client{},
346+
}
347+
348+
if _, err := m.GetSandboxBySession(context.Background(), "", "default", "test-runtime", types.AgentRuntimeKind); err != nil {
349+
t.Fatalf("GetSandboxBySession unexpected error: %v", err)
350+
}
351+
}
352+
353+
func TestGetSandboxBySession_CreateSandbox_TokenFileReadError(t *testing.T) {
354+
dir := t.TempDir()
355+
356+
origTokenPath := serviceAccountTokenPath
357+
serviceAccountTokenPath = dir // ReadFile on a directory returns an error (not IsNotExist)
358+
defer func() {
359+
serviceAccountTokenPath = origTokenPath
360+
}()
361+
362+
_ = os.Unsetenv("API_TOKEN")
363+
364+
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
365+
if got := r.Header.Get("Authorization"); got != "" {
366+
t.Fatalf("expected no Authorization header, got %q", got)
367+
}
368+
resp := types.CreateSandboxResponse{
369+
SessionID: "new-session-123",
370+
SandboxID: "sandbox-456",
371+
SandboxName: "sandbox-test",
372+
EntryPoints: []types.SandboxEntryPoint{{Endpoint: "10.0.0.1:9000"}},
373+
}
374+
w.Header().Set("Content-Type", "application/json")
375+
w.WriteHeader(http.StatusOK)
376+
_ = json.NewEncoder(w).Encode(resp)
377+
}))
378+
defer mockServer.Close()
379+
380+
m := &manager{
381+
storeClient: &fakeStoreClient{},
382+
workloadMgrAddr: mockServer.URL,
383+
httpClient: &http.Client{},
384+
}
385+
386+
if _, err := m.GetSandboxBySession(context.Background(), "", "default", "test-runtime", types.AgentRuntimeKind); err != nil {
387+
t.Fatalf("GetSandboxBySession unexpected error: %v", err)
388+
}
389+
}
390+
229391
func TestGetSandboxBySession_CreateSandbox_CodeInterpreter_Success(t *testing.T) {
230392
// Mock workload manager server
231393
mockServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

0 commit comments

Comments
 (0)