Files
hpc/internal/slurm/jwt.go
2026-04-09 10:31:16 +08:00

63 lines
1.3 KiB
Go

package slurm
import (
"crypto/hmac"
"crypto/sha256"
"encoding/base64"
"encoding/json"
"fmt"
"os"
"time"
)
type JWTClaims struct {
Sun string `json:"sun"`
IAT int64 `json:"iat"`
EXP int64 `json:"exp"`
}
func SignJWT(key []byte, username string, lifespan time.Duration) (string, error) {
if username == "" {
return "", fmt.Errorf("username must not be empty")
}
now := time.Now()
header := map[string]string{"alg": "HS256", "typ": "JWT"}
claims := JWTClaims{
Sun: username,
IAT: now.Unix(),
EXP: now.Add(lifespan).Unix(),
}
headerJSON, err := json.Marshal(header)
if err != nil {
return "", fmt.Errorf("marshal header: %w", err)
}
claimsJSON, err := json.Marshal(claims)
if err != nil {
return "", fmt.Errorf("marshal claims: %w", err)
}
enc := base64.RawURLEncoding
headerEnc := enc.EncodeToString(headerJSON)
claimsEnc := enc.EncodeToString(claimsJSON)
signingInput := headerEnc + "." + claimsEnc
mac := hmac.New(sha256.New, key)
mac.Write([]byte(signingInput))
sig := enc.EncodeToString(mac.Sum(nil))
return signingInput + "." + sig, nil
}
func ReadJWTKey(path string) ([]byte, error) {
data, err := os.ReadFile(path)
if err != nil {
return nil, err
}
if len(data) < 16 {
return nil, fmt.Errorf("key must be at least 16 bytes, got %d", len(data))
}
return data, nil
}