add mock IdP device auth flow

This commit is contained in:
Kenneth Jenkins 2025-04-08 15:59:46 -07:00
parent fe6bb6aea5
commit 039a87dac9
2 changed files with 59 additions and 54 deletions

View file

@ -1,11 +1,7 @@
package ssh
import (
"bytes"
"crypto/ed25519"
"fmt"
"os"
"strings"
"testing"
"github.com/stretchr/testify/require"
@ -36,6 +32,7 @@ func TestSSH(t *testing.T) {
// pomerium + upstream setup
env := testenv.New(t)
env.Add(scenarios.NewIDP([]*scenarios.User{{Email: "test@example.com"}}, scenarios.WithEnableDeviceAuth(true)))
env.Add(scenarios.SSH(scenarios.SSHConfig{}))
env.Add(&ki)
@ -44,13 +41,12 @@ func TestSSH(t *testing.T) {
upstreams.WithAuthorizedKey(clientKey.PublicKey(), "demo"))
r := up.Route().
From(env.SubdomainURLWithScheme("ssh", "ssh")).
Policy(func(p *config.Policy) { p.AllowPublicUnauthenticatedAccess = true })
Policy(func(p *config.Policy) { p.AllowAnyAuthenticatedUser = true })
env.AddUpstream(up)
env.Start()
snippets.WaitStartupComplete(env)
// test scenario -- first verify that the upstream is working at all
//client, err := up.DirectDial(r, clientConfig)
// verify that a connection can be established
client, err := up.Dial(r, clientConfig)
require.NoError(t, err)
defer client.Close()
@ -69,46 +65,3 @@ func newSSHKey(t *testing.T) ssh.Signer {
require.NoError(t, err)
return signer
}
func TestHelloWorld(t *testing.T) {
t.Skip("debugging...")
key, err := os.ReadFile("/Users/kjenkins/scratch/sshd/demo_key")
require.NoError(t, err)
signer, err := ssh.ParsePrivateKey(key)
require.NoError(t, err)
config := &ssh.ClientConfig{
User: "demo",
Auth: []ssh.AuthMethod{
ssh.PublicKeys(signer),
},
HostKeyCallback: ssh.InsecureIgnoreHostKey(),
}
conn, err := ssh.Dial("tcp", "localhost:2222", config)
require.NoError(t, err, "unable to connect")
defer conn.Close()
//conn.ServerVersion()
sess, err := conn.NewSession()
require.NoError(t, err, "unable to start session")
defer sess.Close()
var output bytes.Buffer
sess.Stdout = &output
sess.Stdin = strings.NewReader("whoami\n")
err = sess.Shell()
fmt.Println("Shell() returned ", err)
err = sess.Wait()
fmt.Println("Wait() returned ", err)
fmt.Println(" --> output:\n\n", output.String())
//sess.SendRequest()
}

View file

@ -20,6 +20,8 @@ import (
"strings"
"time"
"golang.org/x/oauth2"
"github.com/go-jose/go-jose/v3"
"github.com/go-jose/go-jose/v3/jwt"
"github.com/google/uuid"
@ -45,7 +47,8 @@ type IDP struct {
}
type IDPOptions struct {
enableTLS bool
enableTLS bool
enableDeviceAuth bool
}
type IDPOption func(*IDPOptions)
@ -62,6 +65,12 @@ func WithEnableTLS(enableTLS bool) IDPOption {
}
}
func WithEnableDeviceAuth(enableDeviceAuth bool) IDPOption {
return func(o *IDPOptions) {
o.enableDeviceAuth = enableDeviceAuth
}
}
// Attach implements testenv.Modifier.
func (idp *IDP) Attach(ctx context.Context) {
env := testenv.EnvFromContext(ctx)
@ -120,7 +129,7 @@ func (idp *IDP) Attach(ctx context.Context) {
router.Handle("/.well-known/openid-configuration", func(w http.ResponseWriter, r *http.Request) {
log.Ctx(ctx).Debug().Str("method", r.Method).Str("uri", r.RequestURI).Send()
rootURL, _ := url.Parse(idp.url.Value())
_ = json.NewEncoder(w).Encode(map[string]interface{}{
config := map[string]interface{}{
"issuer": rootURL.String(),
"authorization_endpoint": rootURL.ResolveReference(&url.URL{Path: "/oidc/auth"}).String(),
"token_endpoint": rootURL.ResolveReference(&url.URL{Path: "/oidc/token"}).String(),
@ -129,11 +138,19 @@ func (idp *IDP) Attach(ctx context.Context) {
"id_token_signing_alg_values_supported": []string{
"ES256",
},
})
}
if idp.enableDeviceAuth {
config["device_authorization_endpoint"] =
rootURL.ResolveReference(&url.URL{Path: "/oidc/device/code"}).String()
}
serveJSON(w, config)
})
router.Handle("/oidc/auth", idp.HandleAuth)
router.Handle("/oidc/token", idp.HandleToken)
router.Handle("/oidc/userinfo", idp.HandleUserInfo)
if idp.enableDeviceAuth {
router.Handle("/oidc/device/code", idp.HandleDeviceCode)
}
env.AddUpstream(router)
}
@ -258,14 +275,25 @@ func (idp *IDP) HandleAuth(w http.ResponseWriter, r *http.Request) {
// HandleToken handles the token flow for OIDC.
func (idp *IDP) HandleToken(w http.ResponseWriter, r *http.Request) {
rawCode := r.FormValue("code")
if idp.enableDeviceAuth && r.FormValue("device_code") != "" {
idp.serveToken(w, r, &State{
ClientID: r.FormValue("client_id"),
Email: "fake.user@example.com",
})
return
}
rawCode := r.FormValue("code")
state, err := DecodeState(rawCode)
if err != nil {
http.Error(w, err.Error(), http.StatusBadRequest)
return
}
idp.serveToken(w, r, state)
}
func (idp *IDP) serveToken(w http.ResponseWriter, r *http.Request, state *State) {
serveJSON(w, map[string]interface{}{
"access_token": state.Encode(),
"refresh_token": state.Encode(),
@ -300,6 +328,30 @@ func (idp *IDP) HandleUserInfo(w http.ResponseWriter, r *http.Request) {
serveJSON(w, state.GetUserInfo(idp.userLookup))
}
// HandleDeviceCode initiates a device auth code flow.
//
// This is the bare minimum to simulate the device auth code flow. There is no client_id
// verification or any actual login.
func (idp *IDP) HandleDeviceCode(w http.ResponseWriter, r *http.Request) {
deviceCode := "GmRhmhcxhwAzkoEqiMEg_DnyEysNkuNhszIySk9eS"
userCode := "ABCD-EFGH"
rootURL, _ := url.Parse(idp.url.Value())
u := rootURL.ResolveReference(&url.URL{Path: "/oidc/device"}) // note: not actually implemented
verificationURI := u.String()
u.RawQuery = "user_code=" + userCode
verificationURIComplete := u.String()
serveJSON(w, &oauth2.DeviceAuthResponse{
DeviceCode: deviceCode,
UserCode: userCode,
VerificationURI: verificationURI,
VerificationURIComplete: verificationURIComplete,
Expiry: time.Now().Add(5 * time.Minute),
Interval: 1,
})
}
type RootURLKey struct{}
var rootURLKey RootURLKey