Refactor with minor improvements, and add README (#5)

.github/workflows/build.yml

@@ -34,7 +34,7 @@ jobs:
         CGO_ENABLED: 0
     - name: Create GitHub Release
       if: startsWith(github.ref, 'refs/tags/')
-      id: upload-release-asset 
+      id: upload-release-asset
       uses: softprops/action-gh-release@v1
       with:
         files: |
@@ -49,7 +49,8 @@ jobs:
         id: meta
         uses: docker/metadata-action@v5
         with:
-          images: ghcr.io/ustc-vlab/sshmux
+          images: |
+            ghcr.io/${{ github.repository_owner }}/sshmux
           tags: |
             type=ref,event=branch
             type=ref,event=pr

README.md

@@ -0,0 +1,72 @@
+# sshmux
+
+`sshmux` is a new, simple implementation of SSH reverse proxy. `sshmux` was initially developed for Vlab, while we'd like to expand its usage to cover more scenarios.
+
+## Build, Run and Test
+
+`sshmux` requires a Go 1.21+ toolchain to build. You can use `go build` or `make` to get the `sshmux` binary directly in the directory.
+
+You can run the built binary with `./sshmux`. Note that you'll need to provide a valid configuration file as described [here](#config).
+
+You can perform unit tests with `go test` or `make test`. Enable verbose logging with `go test -v`.
+
+## Config
+
+`sshmux` requires a JSON configuration file to start up. By default it will look at `/etc/sshmux/config.json`, but you can also specify a custom configuration by passing `-c path/to/config.json` in the command line arguments. An [example](config.example.json) file is provided.
+
+The table below shows the available options for `sshmux`:
+
+| Key         | Type       | Description                                                        | Required | Example                            |
+|-------------|------------|--------------------------------------------------------------------|----------|------------------------------------|
+| `address`   | `string`   | TCP host and port that `sshmux` will listen on.                    | `true`   | `"0.0.0.0:8022"`                   |
+| `host-keys` | `[]string` | Paths to SSH host key files with which `sshmux` identifies itself. | `true`   | `["/sshmux/ssh_host_ed25519_key"]` |
+| `api`       | `string`   | HTTP address that `sshmux` shall interact with.                    | `true`   | `"http://127.0.0.1:5000/ssh"`      |
+| `token`     | `string`   | Token used to authenticate with the API endpoint.                  | `true`   | `"long-and-random-token"`          |
+| `banner`    | `string`   | SSH banner to send to downstream.                                  | `false`  | `"Welcome to Vlab\n"`              |
+| `logger`    | `string`   | UDP host and port that `sshmux` send log messages to.              | `false`  | `"127.0.0.1:5556"`                 |
+| `proxy-protocol-allowed-cidrs` | `[]string` | CIDRs from which [PROXY protocol](https://www.haproxy.com/blog/use-the-proxy-protocol-to-preserve-a-clients-ip-address) is allowed. | `false` | `["127.0.0.22/32"]` |
+
+### Advanced Config
+
+The table below shows extra options for `sshmux`, mainly for authentication with Vlab backends:
+
+| Key                        | Type       | Description                                                                | Example                      |
+|----------------------------|------------|----------------------------------------------------------------------------|------------------------------|
+| `recovery-token`           | `string`   | Token used to authenticate with the recovery backend. Defaults to `token`. | `"long-and-random-token"`    |
+| `recovery-server`          | `string`   | SSH host and port of the recovery server.                                  | `"172.30.0.101:2222"`        |
+| `recovery-username`        | `[]string` | Usernames dedicated to the recovery server.                                | `["recovery", "console"]`    |
+| `all-username-nopassword`  | `bool`     | If set to `true`, no users will be asked for UNIX password.                | `true`                       |
+| `username-nopassword`      | `[]string` | Usernames that won't be asked for UNIX password.                           | `["vlab", "ubuntu", "root"]` |
+| `invalid-username`         | `[]string` | Usernames that are known to be invalid.                                    | `["user"]`                   |
+| `invalid-username-message` | `string`   | Message to display when the requested username is invalid.                 | `"Invalid username %s."`     |
+
+All of these options can be omitted, if the corresponding feature is not intended to be used.
+
+## API server
+
+`sshmux` requires an API server to perform authentication and authorization for a user.
+
+The API accepts JSON input with the following keys:
+
+| Key               | Type     | Description                                                                                              |
+|-------------------|----------|----------------------------------------------------------------------------------------------------------|
+| `auth_type`       | `string` | The authentication type. Always set to `"key"` at the moment.                                            |
+| `username`        | `string` | Vlab username. Omitted if the user is authenticating with public key.                                    |
+| `password`        | `string` | Vlab password. Omitted if the user is authenticating with public key.                                    |
+| `public_key_type` | `string` | SSH public key type. Omitted if the user is authenticating with username and password.                   |
+| `public_key_data` | `string` | Base64-encoded SSH public key payload. Omitted if the user is authenticating with username and password. |
+| `unix_username`   | `string` | UNIX username the user is requesting access to.                                                          |
+| `token`           | `string` | Token used to authenticate the `sshmux` instance.                                                        |
+
+The API responds with JSON output with the following keys:
+
+| Key              | Type      | Description                                                                                                          |
+|------------------|-----------|----------------------------------------------------------------------------------------------------------------------|
+| `status`         | `string`  | The authentication status. Should be `"ok"` if the user is authorized.                                               |
+| `address`        | `string`  | TCP host and port of the downstream SSH server the user is requesting for.                                           |
+| `private_key`    | `string`  | SSH private key to authenticate for the downstream.                                                                  |
+| `cert`           | `string`  | The certificate associated with the SSH private key.                                                                 |
+| `vmid`           | `integer` | ID of the requested VM. Only used for recovery access.                                                               |
+| `proxy_protocol` | `integer` | PROXY protocol version to use for the downstream. Should be `1`, `2` or omitted (which disables PROXY protocol). |
+
+Note that if the user is not authorized, the API server should return a `status` other than `"ok"`, and other keys can be safely ommitted.

auth.go

@@ -0,0 +1,160 @@
+package main
+
+import (
+	"bytes"
+	"encoding/base64"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"slices"
+
+	"golang.org/x/crypto/ssh"
+)
+
+type AuthRequestPublicKey struct {
+	AuthType      string `json:"auth_type"`
+	UnixUsername  string `json:"unix_username"`
+	PublicKeyType string `json:"public_key_type"`
+	PublicKeyData string `json:"public_key_data"`
+	Token         string `json:"token"`
+}
+
+type AuthRequestPassword struct {
+	AuthType     string `json:"auth_type"`
+	Username     string `json:"username"`
+	Password     string `json:"password"`
+	UnixUsername string `json:"unix_username"`
+	Token        string `json:"token"`
+}
+
+type AuthResponse struct {
+	Status        string `json:"status"`
+	Address       string `json:"address"`
+	PrivateKey    string `json:"private_key"`
+	Cert          string `json:"cert"`
+	Id            int    `json:"vmid"`
+	ProxyProtocol byte   `json:"proxy_protocol,omitempty"`
+}
+
+type UpstreamInformation struct {
+	Host          string
+	Signer        ssh.Signer
+	Password      *string
+	ProxyProtocol byte
+}
+
+type Authenticator struct {
+	Endpoint string
+	Token    string
+	Recovery RecoveryConfig
+}
+
+func makeAuthenticator(config Config) Authenticator {
+	recoveryToken := config.RecoveryToken
+	if recoveryToken == "" {
+		recoveryToken = config.Token
+	}
+	return Authenticator{
+		Endpoint: config.API,
+		Token:    config.Token,
+		Recovery: RecoveryConfig{
+			Server:   config.RecoveryServer,
+			Username: config.RecoveryUsername,
+			Token:    recoveryToken,
+		},
+	}
+}
+
+func parsePrivateKey(key string, cert string) ssh.Signer {
+	if key == "" {
+		return nil
+	}
+	signer, err := ssh.ParsePrivateKey([]byte(key))
+	if err != nil {
+		return nil
+	}
+	if cert == "" {
+		return signer
+	}
+	pk, _, _, _, err := ssh.ParseAuthorizedKey([]byte(cert))
+	if err != nil {
+		return signer
+	}
+	certSigner, err := ssh.NewCertSigner(pk.(*ssh.Certificate), signer)
+	if err != nil {
+		return signer
+	}
+	return certSigner
+}
+
+func (auth Authenticator) AuthUser(request any, username string) (*UpstreamInformation, error) {
+	payload := new(bytes.Buffer)
+	if err := json.NewEncoder(payload).Encode(request); err != nil {
+		return nil, err
+	}
+	res, err := http.Post(auth.Endpoint, "application/json", payload)
+	if err != nil {
+		return nil, err
+	}
+	defer res.Body.Close()
+	body, err := io.ReadAll(res.Body)
+	if err != nil {
+		return nil, err
+	}
+	var response AuthResponse
+	err = json.Unmarshal(body, &response)
+	if err != nil {
+		return nil, err
+	}
+	if response.Status != "ok" {
+		return nil, nil
+	}
+
+	var upstream UpstreamInformation
+	// FIXME: Can this be handled in API server?
+	if slices.Contains(auth.Recovery.Username, username) {
+		upstream.Host = auth.Recovery.Server
+		password := fmt.Sprintf("%d %s", response.Id, auth.Recovery.Token)
+		upstream.Password = &password
+	} else {
+		upstream.Host = response.Address
+	}
+	upstream.Signer = parsePrivateKey(response.PrivateKey, response.Cert)
+	upstream.ProxyProtocol = response.ProxyProtocol
+	return &upstream, nil
+}
+
+func (auth Authenticator) AuthUserWithPublicKey(key ssh.PublicKey, unixUsername string) (*UpstreamInformation, error) {
+	keyType := key.Type()
+	keyData := base64.StdEncoding.EncodeToString(key.Marshal())
+	request := &AuthRequestPublicKey{
+		AuthType:      "key",
+		UnixUsername:  unixUsername,
+		PublicKeyType: keyType,
+		PublicKeyData: keyData,
+		Token:         auth.Token,
+	}
+	return auth.AuthUser(request, unixUsername)
+}
+
+func (auth Authenticator) AuthUserWithUserPass(username string, password string, unixUsername string) (*UpstreamInformation, error) {
+	request := &AuthRequestPassword{
+		AuthType:     "key",
+		Username:     username,
+		Password:     password,
+		UnixUsername: unixUsername,
+		Token:        auth.Token,
+	}
+	return auth.AuthUser(request, unixUsername)
+}
+
+func removePublicKeyMethod(methods []string) []string {
+	res := make([]string, 0, len(methods))
+	for _, s := range methods {
+		if s != "publickey" {
+			res = append(res, s)
+		}
+	}
+	return res
+}

config.go

@@ -0,0 +1,35 @@
+package main
+
+type Config struct {
+	Address    string   `json:"address"`
+	ProxyCIDRs []string `json:"proxy-protocol-allowed-cidrs"`
+	HostKeys   []string `json:"host-keys"`
+	API        string   `json:"api"`
+	Logger     string   `json:"logger"`
+	Banner     string   `json:"banner"`
+	Token      string   `json:"token"`
+	// The following should be moved into API server
+	RecoveryToken          string   `json:"recovery-token"`
+	RecoveryServer         string   `json:"recovery-server"`
+	RecoveryUsername       []string `json:"recovery-username"`
+	AllUsernameNoPassword  bool     `json:"all-username-nopassword"`
+	UsernameNoPassword     []string `json:"username-nopassword"`
+	InvalidUsername        []string `json:"invalid-username"`
+	InvalidUsernameMessage string   `json:"invalid-username-message"`
+}
+
+type UsernamePolicyConfig struct {
+	InvalidUsername        []string `json:"invalid-username"`
+	InvalidUsernameMessage string   `json:"invalid-username-message"`
+}
+
+type PasswordPolicyConfig struct {
+	AllUsernameNoPassword bool     `json:"all-username-nopassword"`
+	UsernameNoPassword    []string `json:"username-nopassword"`
+}
+
+type RecoveryConfig struct {
+	Server   string   `json:"recovery-server"`
+	Username []string `json:"recovery-username"`
+	Token    string   `json:"token"`
+}

logging.go

@@ -0,0 +1,52 @@
+package main
+
+import (
+	"encoding/json"
+	"log"
+	"net"
+	"time"
+)
+
+type LogMessage struct {
+	LoginTime      int64  `json:"login_time"`
+	DisconnectTime int64  `json:"disconnect_time"`
+	ClientIp       string `json:"remote_ip"`
+	HostIp         string `json:"host_ip"`
+	Username       string `json:"user_name"`
+}
+
+type Logger struct {
+	Channel chan LogMessage
+}
+
+func makeLogger(url string) Logger {
+	channel := make(chan LogMessage, 256)
+	go func() {
+		if url == "" {
+			for range channel {
+			}
+			return
+		}
+		conn, err := net.Dial("udp", url)
+		if err != nil {
+			log.Printf("Logger Dial failed: %s\n", err)
+			// Drain the channel to avoid blocking
+			for range channel {
+			}
+			return
+		}
+		for logMessage := range channel {
+			jsonMsg, err := json.Marshal(logMessage)
+			if err != nil {
+				continue
+			}
+			conn.Write(jsonMsg)
+		}
+	}()
+	return Logger{Channel: channel}
+}
+
+func (l Logger) SendLog(logMessage *LogMessage) {
+	logMessage.DisconnectTime = time.Now().Unix()
+	l.Channel <- *logMessage
+}

main.go

@@ -0,0 +1,32 @@
+package main
+
+import (
+	"encoding/json"
+	"flag"
+	"log"
+	"os"
+)
+
+func sshmuxServer(configFile string) {
+	var config Config
+	configFileBytes, err := os.ReadFile(configFile)
+	if err != nil {
+		log.Fatal(err)
+	}
+	err = json.Unmarshal(configFileBytes, &config)
+	if err != nil {
+		log.Fatal(err)
+	}
+	sshmux, err := makeServer(config)
+	if err != nil {
+		log.Fatal(err)
+	}
+	sshmux.ListenAddr(config.Address)
+}
+
+func main() {
+	var configFile string
+	flag.StringVar(&configFile, "c", "/etc/sshmux/config.json", "config file")
+	flag.Parse()
+	sshmuxServer(configFile)
+}