package main
import (
"bytes"
"context"
"encoding/json"
"fmt"
"log"
"net"
"net/http"
"time"
"github.com/thuongh2/go-security/agentauth"
"github.com/thuongh2/go-security/jwt"
"github.com/thuongh2/go-security/oauth2"
)
const agentPrincipalID = "agent-key-abc123" // stand-in for an existing API key ID / OAuth2 client ID
func main() {
listener, err := net.Listen("tcp", "127.0.0.1:0")
if err != nil {
log.Fatalf("listen: %v", err)
}
baseURL := fmt.Sprintf("http://%s", listener.Addr())
executeURL := baseURL + "/capability/execute"
// --- demo signing key (HS256, shared secret) ---
// PRODUCTION: use an asymmetric key (RS256/ES256/EdDSA) so the token
// verifier never holds a signing secret.
key := jwt.NewHS256([]byte("demo-only-hs256-secret-change-me-32bytes!!"))
signer, err := jwt.NewSigner(key, "demo-key")
if err != nil {
log.Fatalf("signer: %v", err)
}
verifier, err := jwt.NewVerifierKey(jwt.HS256, key)
if err != nil {
log.Fatalf("verifier: %v", err)
}
tokenCfg := agentauth.TokenConfig{
Signer: signer,
Verifier: verifier,
Issuer: baseURL,
Audience: executeURL,
TTL: 2 * time.Second, // deliberately short for this demo
Leeway: time.Minute,
}
// --- agentauth primitives ---
registry := agentauth.NewRegistry()
if err := registry.Register(agentauth.Capability{
Name: "greet",
Description: "Return a greeting for the caller.",
Handler: func(_ context.Context, sess agentauth.AgentSession, _ json.RawMessage) (any, error) {
return map[string]string{"message": "hello, " + sess.PrincipalID}, nil
},
}); err != nil {
log.Fatalf("register capability: %v", err)
}
grants := agentauth.NewMemoryGrantStore()
replay := agentauth.NewMemoryReplayStore()
pending := agentauth.NewMemoryPendingRequestStore()
// --- oauth2 device-authorization flow ---
deviceStore := oauth2.NewMemoryDeviceCodeStore()
deviceSvc := oauth2.NewDeviceAuthorizationService(deviceStore, baseURL+"/device/verify", 10*time.Minute, 0)
issueToken := agentauth.NewAccessTokenIssuer(tokenCfg, grants, pending, 24*time.Hour)
// --- HTTP wiring ---
mux := http.NewServeMux()
mux.Handle("/.well-known/agent-configuration", agentauth.DiscoveryHandler(agentauth.DiscoveryConfig{
Issuer: baseURL,
ProviderName: "go-security agentauth demo",
ExecuteURL: executeURL,
CapabilitiesURL: baseURL + "/capability/list",
Modes: []string{"delegated"},
ApprovalEndpoints: map[string]string{
"device_authorization": baseURL + "/device/authorize",
},
}))
mux.Handle("/capability/list", agentauth.CapabilitiesHandler(registry))
mux.Handle("/capability/execute", agentauth.ExecuteHandler(agentauth.ExecuteDeps{
Registry: registry,
Grants: grants,
Replay: replay,
Token: tokenCfg,
}))
mux.Handle("/device/token", oauth2.DeviceTokenHandler(deviceSvc, issueToken))
srv := &http.Server{Handler: mux}
go func() {
if err := srv.Serve(listener); err != nil && err != http.ErrServerClosed {
log.Fatalf("serve: %v", err)
}
}()
defer srv.Close()
client := &http.Client{Timeout: 5 * time.Second}
// --- Step 1: agent begins the device flow with its own credential ---
auth, err := deviceSvc.Begin(context.Background())
if err != nil {
log.Fatalf("begin: %v", err)
}
if err := pending.Create(context.Background(), auth.UserCode, agentPrincipalID, []string{"greet"}, 10*time.Minute); err != nil {
log.Fatalf("pending.Create: %v", err)
}
fmt.Printf("1. agent began device flow: user_code=%s\n", auth.UserCode)
// --- Step 2: the app's CUSTOM approve handler (not DeviceApproveHandler) ---
// simulates the human clicking "approve" on the verification page.
req, err := pending.GetByUserCode(context.Background(), auth.UserCode)
if err != nil || req == nil {
log.Fatalf("pending.GetByUserCode: %v (req=%v)", err, req)
}
if err := deviceSvc.Authorize(context.Background(), auth.UserCode, req.AgentPrincipalID); err != nil {
log.Fatalf("authorize: %v", err)
}
fmt.Printf("2. human approved for agent principal=%s, capabilities=%v\n", req.AgentPrincipalID, req.RequestedCapabilities)
// --- Step 3: agent polls once and gets its first capability token ---
tok := pollForToken(client, baseURL, auth.DeviceCode)
fmt.Printf("3. agent obtained capability token (expires in %ds)\n", int(tokenCfg.TTL.Seconds()))
// --- Step 4: agent calls the capability ---
callGreet(client, executeURL, tok)
// --- Step 5: token expires; no device-flow "refresh" exists (Poll is
// single-use) — the app mints a fresh token directly off the still-Active
// grant via RefreshCapabilityToken, simulating its own authenticated
// internal channel (never exposed as a bare network endpoint).
fmt.Println("5. waiting for the capability token to expire...")
time.Sleep(tokenCfg.TTL + 500*time.Millisecond)
refreshed, err := agentauth.RefreshCapabilityToken(context.Background(), agentPrincipalID, grants, tokenCfg)
if err != nil {
log.Fatalf("refresh: %v", err)
}
fmt.Println(" refreshed token off the same 24h grant (no second approval needed)")
callGreet(client, executeURL, refreshed)
}
func pollForToken(client *http.Client, baseURL, deviceCode string) string {
form := "device_code=" + deviceCode
resp, err := client.Post(baseURL+"/device/token", "application/x-www-form-urlencoded", bytes.NewBufferString(form))
if err != nil {
log.Fatalf("poll: %v", err)
}
defer resp.Body.Close()
var tr oauth2.TokenResponse
if err := json.NewDecoder(resp.Body).Decode(&tr); err != nil {
log.Fatalf("decode token response: %v", err)
}
if tr.AccessToken == "" {
log.Fatalf("poll: empty access token (status %d)", resp.StatusCode)
}
return tr.AccessToken
}
func callGreet(client *http.Client, executeURL, token string) {
body, _ := json.Marshal(map[string]any{"capability": "greet", "arguments": map[string]string{}})
req, err := http.NewRequest(http.MethodPost, executeURL, bytes.NewReader(body))
if err != nil {
log.Fatalf("new request: %v", err)
}
req.Header.Set("Authorization", "Bearer "+token)
resp, err := client.Do(req)
if err != nil {
log.Fatalf("execute: %v", err)
}
defer resp.Body.Close()
var out map[string]string
if err := json.NewDecoder(resp.Body).Decode(&out); err != nil {
log.Fatalf("decode execute response: %v", err)
}
fmt.Printf(" -> capability result (status %d): %v\n", resp.StatusCode, out)
}
