SimpleRemoter/server/go/web/sso_test.go

package web

import (
	"encoding/base64"
	"encoding/json"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/golang-jwt/jwt/v5"

	"github.com/yuanyuanxiang/SimpleRemoter/server/go/logger"
	"github.com/yuanyuanxiang/SimpleRemoter/server/go/wsauth"
)

const testSSOSecret = "test-shared-secret"

func signTestToken(t *testing.T, secret string, claims jwt.MapClaims) string {
	t.Helper()
	tok := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
	s, err := tok.SignedString([]byte(secret))
	if err != nil {
		t.Fatalf("sign test token: %v", err)
	}
	return s
}

// noneToken builds a base64url header/claims JWT with alg "none" and no
// signature, to exercise the algorithm-confusion rejection path.
func noneToken(claims jwt.MapClaims) string {
	h, _ := json.Marshal(map[string]string{"alg": "none", "typ": "JWT"})
	c, _ := json.Marshal(claims)
	enc := base64.RawURLEncoding.EncodeToString
	return strings.TrimRight(enc(h), "=") + "." + strings.TrimRight(enc(c), "=") + "."
}

func cookieReq(token string) *http.Request {
	r := &http.Request{Header: http.Header{}}
	r.AddCookie(&http.Cookie{Name: "auth_token", Value: token})
	return r
}

func bearerReq(token string) *http.Request {
	r := &http.Request{Header: http.Header{}}
	r.Header.Set("Authorization", "Bearer "+token)
	return r
}

func TestJWTDetector(t *testing.T) {
	d := NewJWTDetector(JWTDetectorConfig{Secret: testSSOSecret})

	valid := func(rolename string) jwt.MapClaims {
		return jwt.MapClaims{
			"username": "alice",
			"rolename": rolename,
			"exp":      time.Now().Add(time.Hour).Unix(),
		}
	}

	tests := []struct {
		name     string
		req      *http.Request
		wantUser string
		wantRole string
		wantOK   bool
	}{
		{"admin maps to admin", cookieReq(signTestToken(t, testSSOSecret, valid("admin"))), "alice", "admin", true},
		{"viewer maps to viewer", cookieReq(signTestToken(t, testSSOSecret, valid("viewer"))), "alice", "viewer", true},
		{"empty rolename maps to viewer", cookieReq(signTestToken(t, testSSOSecret, valid(""))), "alice", "viewer", true},
		{"unknown rolename maps to viewer", cookieReq(signTestToken(t, testSSOSecret, valid("operator"))), "alice", "viewer", true},
		{"bearer header accepted", bearerReq(signTestToken(t, testSSOSecret, valid("admin"))), "alice", "admin", true},
		{"no credential", &http.Request{Header: http.Header{}}, "", "", false},
		{"wrong secret rejected", cookieReq(signTestToken(t, "other-secret", valid("admin"))), "", "", false},
		{"expired token rejected", cookieReq(signTestToken(t, testSSOSecret, jwt.MapClaims{
			"username": "alice", "rolename": "admin", "exp": time.Now().Add(-time.Hour).Unix(),
		})), "", "", false},
		{"missing username rejected", cookieReq(signTestToken(t, testSSOSecret, jwt.MapClaims{
			"rolename": "admin", "exp": time.Now().Add(time.Hour).Unix(),
		})), "", "", false},
		{"alg none rejected", cookieReq(noneToken(valid("admin"))), "", "", false},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			user, role, ok := d.Detect(tt.req)
			if ok != tt.wantOK {
				t.Fatalf("ok = %v, want %v (user=%q role=%q)", ok, tt.wantOK, user, role)
			}
			if !tt.wantOK {
				return
			}
			if user != tt.wantUser || role != tt.wantRole {
				t.Fatalf("got (%q, %q), want (%q, %q)", user, role, tt.wantUser, tt.wantRole)
			}
		})
	}
}

// A detector with no shared secret must never accept a token, even one signed
// with an empty key — otherwise a misconfigured deployment would silently
// accept forged sessions.
func TestJWTDetectorEmptySecretFailClosed(t *testing.T) {
	d := NewJWTDetector(JWTDetectorConfig{})
	tok := signTestToken(t, "", jwt.MapClaims{
		"username": "alice", "rolename": "admin", "exp": time.Now().Add(time.Hour).Unix(),
	})
	if _, _, ok := d.Detect(cookieReq(tok)); ok {
		t.Fatal("empty-secret detector accepted a token; expected fail-closed")
	}
}

// Non-default config keys must be honored, so a different portal can reuse the
// detector with no code changes.
func TestJWTDetectorCustomConfig(t *testing.T) {
	d := NewJWTDetector(JWTDetectorConfig{
		Secret:        testSSOSecret,
		CookieName:    "session",
		UsernameClaim: "sub",
		RoleClaim:     "role",
		AdminRole:     "root",
	})
	tok := signTestToken(t, testSSOSecret, jwt.MapClaims{
		"sub": "bob", "role": "root", "exp": time.Now().Add(time.Hour).Unix(),
	})
	r := &http.Request{Header: http.Header{}}
	r.AddCookie(&http.Cookie{Name: "session", Value: tok})
	user, role, ok := d.Detect(r)
	if !ok || user != "bob" || role != "admin" {
		t.Fatalf("got (%q, %q, %v), want (%q, %q, true)", user, role, ok, "bob", "admin")
	}
}

// fakeSSODetector lets the /sso handler's role policy be exercised without a
// real portal JWT.
type fakeSSODetector struct {
	username, role string
	ok             bool
}

func (f *fakeSSODetector) Detect(*http.Request) (string, string, bool) {
	return f.username, f.role, f.ok
}

// The /sso handler must only bridge administrators. A valid portal session for
// any other role — or no session at all — returns 401 so the front-end falls
// back to the built-in login form.
func TestHandleSSOAdminOnly(t *testing.T) {
	newSrv := func(d SSODetector) *Server {
		return &Server{
			log:  logger.New(logger.Config{}),
			auth: wsauth.New(),
			sso:  d,
		}
	}

	tests := []struct {
		name       string
		detector   SSODetector
		wantStatus int
	}{
		{"admin is bridged", &fakeSSODetector{"alice", "admin", true}, http.StatusOK},
		{"viewer is rejected", &fakeSSODetector{"bob", "viewer", true}, http.StatusUnauthorized},
		{"unknown role is rejected", &fakeSSODetector{"carol", "operator", true}, http.StatusUnauthorized},
		{"no credential is rejected", &fakeSSODetector{ok: false}, http.StatusUnauthorized},
		{"SSO disabled", nil, http.StatusUnauthorized},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			s := newSrv(tt.detector)
			rec := httptest.NewRecorder()
			s.handleSSO(rec, httptest.NewRequest(http.MethodPost, "/sso", nil))
			if rec.Code != tt.wantStatus {
				t.Fatalf("status = %d, want %d (body=%s)", rec.Code, tt.wantStatus, rec.Body.String())
			}
		})
	}
}