From 928eb1236e4eb5f8ab1d4f9f2290b7c3672ecd97 Mon Sep 17 00:00:00 2001
From: YR Chen <yrchen@mail.ustc.edu.cn>
Date: Sat, 17 Aug 2024 02:24:09 +0800
Subject: [PATCH] Simplify API

---
 auth.go   | 2 +-
 config.go | 8 ++------
 main.go   | 2 +-
 sshmux.go | 6 ++----
 4 files changed, 6 insertions(+), 12 deletions(-)

diff --git a/auth.go b/auth.go
index 33a715f..5fdba15 100644
--- a/auth.go
+++ b/auth.go
@@ -106,7 +106,7 @@ func (auth Authenticator) AuthUser(request any, username string) (*UpstreamInfor
 	var upstream UpstreamInformation
 	// FIXME: Can this be handled in API server?
 	if slices.Contains(auth.Recovery.Usernames, username) {
-		upstream.Host = auth.Recovery.Server
+		upstream.Host = auth.Recovery.Address
 		password := fmt.Sprintf("%d %s", response.Id, auth.Recovery.Token)
 		upstream.Password = &password
 	} else {
diff --git a/config.go b/config.go
index 0329672..374bd36 100644
--- a/config.go
+++ b/config.go
@@ -1,9 +1,5 @@
 package main
 
-type ServerConfig struct {
-	Address string `toml:"address"`
-}
-
 type SSHConfig struct {
 	Banner   string   `toml:"banner"`
 	HostKeys []string `toml:"host-keys"`
@@ -30,13 +26,13 @@ type ProxyProtocolConfig struct {
 }
 
 type RecoveryConfig struct {
-	Server    string   `toml:"server"`
+	Address   string   `toml:"address"`
 	Usernames []string `toml:"usernames"`
 	Token     string   `toml:"token"`
 }
 
 type Config struct {
-	Server        ServerConfig        `toml:"server"`
+	Address       string              `toml:"address"`
 	SSH           SSHConfig           `toml:"ssh"`
 	Auth          AuthConfig          `toml:"auth"`
 	Logger        LoggerConfig        `toml:"logger"`
diff --git a/main.go b/main.go
index f98d333..3707f8e 100644
--- a/main.go
+++ b/main.go
@@ -25,7 +25,7 @@ func sshmuxServer(configFile string) {
 		if err != nil {
 			log.Fatal(err)
 		}
-		sshmux.ListenAddr(config.Server.Address)
+		sshmux.ListenAddr(config.Address)
 	} else {
 		var legacyConfig LegacyConfig
 		configFileBytes, err := os.ReadFile(configFile)
diff --git a/sshmux.go b/sshmux.go
index ffe9e9e..591bbb6 100644
--- a/sshmux.go
+++ b/sshmux.go
@@ -76,9 +76,7 @@ func makeLegacyServer(config LegacyConfig) (*Server, error) {
 		config.RecoveryToken = config.Token
 	}
 	return makeServer(Config{
-		Server: ServerConfig{
-			Address: config.Address,
-		},
+		Address: config.Address,
 		SSH: SSHConfig{
 			Banner:   config.Banner,
 			HostKeys: config.HostKeys,
@@ -100,7 +98,7 @@ func makeLegacyServer(config LegacyConfig) (*Server, error) {
 			Networks: config.ProxyCIDRs,
 		},
 		Recovery: RecoveryConfig{
-			Server:    config.RecoveryServer,
+			Address:   config.RecoveryServer,
 			Usernames: config.RecoveryUsername,
 			Token:     config.RecoveryToken,
 		},