x/crypto/ssh: verified public key callback and arbitrary data in Permissions

// This program demonstrates the intended behaviour of the new
// `VerifiedPublicKeyCallback` and the `ExtraData` field in `Permissions`.
// Since the change is not yet in Go 1.27.1, the code below uses mock
// types that mimic the relevant parts of golang.org/x/crypto/ssh.
// Running this program prints the sequence of callbacks and the
// data passed between them.

package main

import (
	"fmt"
)

// Mock types that resemble the real ssh package structures.
type ConnMetadata struct{ ID int }
type PublicKey struct{ Key string }
type Permissions struct {
	CriticalOptions map[string]string
	Extensions      map[string]string
	ExtraData       any
}
type PartialSuccessError struct{ msg string }
func (e *PartialSuccessError) Error() string { return e.msg }

// ServerConfig holds the callbacks.
type ServerConfig struct {
	PublicKeyCallback       func(ConnMetadata, PublicKey) (*Permissions, error)
	VerifiedPublicKeyCallback func(ConnMetadata, PublicKey, *Permissions) (*Permissions, error)
}

// Simulate an SSH handshake that triggers the callbacks.
func simulateHandshake(cfg ServerConfig, meta ConnMetadata, key PublicKey) {
	// Step 1: PublicKeyCallback is called.
	perms, err := cfg.PublicKeyCallback(meta, key)
	if err != nil {
		fmt.Println("PublicKeyCallback error:", err)
		return
	}
	fmt.Println("PublicKeyCallback returned permissions:", perms)

	// Step 2: If a signature is provided (simulated here), call VerifiedPublicKeyCallback.
	if cfg.VerifiedPublicKeyCallback != nil {
		newPerms, err := cfg.VerifiedPublicKeyCallback(meta, key, perms)
		if err != nil {
			fmt.Println("VerifiedPublicKeyCallback error:", err)
			return
		}
		fmt.Println("VerifiedPublicKeyCallback returned permissions:", newPerms)
	}
}

func main() {
	// Define the callbacks.
	cfg := ServerConfig{
		PublicKeyCallback: func(meta ConnMetadata, key PublicKey) (*Permissions, error) {
			// Pretend we decide the key is acceptable and attach some data.
			perms := &Permissions{
				CriticalOptions: map[string]string{"no-port-forwarding": ""},
				Extensions:      map[string]string{"permit-agent-forwarding": ""},
				ExtraData:       fmt.Sprintf("user-%d", meta.ID),
			}
			return perms, nil
		},
		VerifiedPublicKeyCallback: func(meta ConnMetadata, key PublicKey, perms *Permissions) (*Permissions, error) {
			// Use the data from PublicKeyCallback and modify permissions.
			if perms.ExtraData == nil {
				return nil, &PartialSuccessError{"missing extra data"}
			}
			// Append a new critical option based on the extra data.
			perms.CriticalOptions["permit-pty"] = ""
			return perms, nil
		},
	}

	// Run the simulation.
	meta := ConnMetadata{ID: 42}
	key := PublicKey{Key: "ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABAQC..."}
	simulateHandshake(cfg, meta, key)
}