|
|
|
@ -6,32 +6,32 @@ import ( |
|
|
|
"encoding/json" |
|
|
|
"errors" |
|
|
|
"fmt" |
|
|
|
"github.com/redis/go-redis/v9" |
|
|
|
"log" |
|
|
|
"memobus_relay_server/config" |
|
|
|
"memobus_relay_server/registry" |
|
|
|
"memobus_relay_server/storage" |
|
|
|
"net" |
|
|
|
"net/http" |
|
|
|
"net/http/httputil" |
|
|
|
"os" |
|
|
|
"os/signal" |
|
|
|
"strings" |
|
|
|
"sync" |
|
|
|
"syscall" |
|
|
|
"time" |
|
|
|
|
|
|
|
"github.com/golang-jwt/jwt/v5" |
|
|
|
"github.com/hashicorp/yamux" |
|
|
|
) |
|
|
|
|
|
|
|
// ==============================================================================
|
|
|
|
// 1. 密钥配置
|
|
|
|
// ==============================================================================
|
|
|
|
var ( |
|
|
|
// 用于验证 App 请求的密钥，必须和 ibserver 的 AppAccessSecret 一致
|
|
|
|
appAccessSecret = []byte(os.Getenv("APP_ACCESS_SECRET")) |
|
|
|
// 用于验证设备连接的密钥，必须和旧中继服务的 RelaySecret 一致
|
|
|
|
deviceRelaySecret = []byte(os.Getenv("RELAY_SECRET")) |
|
|
|
appAccessSecret []byte |
|
|
|
deviceRelaySecret []byte |
|
|
|
) |
|
|
|
|
|
|
|
// ==============================================================================
|
|
|
|
// 2. 结构体定义
|
|
|
|
// ==============================================================================
|
|
|
|
type AuthPayload struct { |
|
|
|
DeviceSN string `json:"device_sn"` |
|
|
|
Token string `json:"token"` |
|
|
|
@ -53,20 +53,62 @@ var ( |
|
|
|
sessionMutex = &sync.RWMutex{} |
|
|
|
) |
|
|
|
|
|
|
|
// ==============================================================================
|
|
|
|
// 3. Main & 服务器启动逻辑
|
|
|
|
// ==============================================================================
|
|
|
|
func main() { |
|
|
|
if len(appAccessSecret) == 0 || len(deviceRelaySecret) == 0 { |
|
|
|
log.Println("WARNING: APP_ACCESS_SECRET or RELAY_SECRET environment variable not set.") |
|
|
|
// 1. 加载配置
|
|
|
|
config.LoadConfig() |
|
|
|
appAccessSecret = []byte(config.Cfg.Auth.AppAccessSecret) |
|
|
|
deviceRelaySecret = []byte(config.Cfg.Auth.DeviceRelaySecret) |
|
|
|
|
|
|
|
// 2. 初始化存储层 (Redis)
|
|
|
|
if err := storage.InitRedis(); err != nil { |
|
|
|
log.Fatalf("Failed to initialize storage: %v", err) |
|
|
|
} |
|
|
|
go listenForDevices(":7002") |
|
|
|
|
|
|
|
log.Println("Starting App HTTP server on :8089") |
|
|
|
http.HandleFunc("/tunnel/", handleAppRequest) // 统一入口
|
|
|
|
if err := http.ListenAndServe(":8089", nil); err != nil { |
|
|
|
log.Fatalf("Failed to start App server: %v", err) |
|
|
|
// 3. 启动注册与心跳 (它会自己检查 Redis 是否启用)
|
|
|
|
registry.StartHeartbeat(func() int { |
|
|
|
sessionMutex.RLock() |
|
|
|
defer sessionMutex.RUnlock() |
|
|
|
return len(deviceSessions) |
|
|
|
}) |
|
|
|
|
|
|
|
// 4. 启动核心服务 (放入后台 goroutine)
|
|
|
|
go listenForDevices(config.Cfg.Server.DeviceListenPort) |
|
|
|
|
|
|
|
mux := http.NewServeMux() |
|
|
|
mux.HandleFunc("/tunnel/", handleAppRequest) |
|
|
|
httpServer := &http.Server{ |
|
|
|
Addr: config.Cfg.Server.AppListenPort, |
|
|
|
Handler: mux, |
|
|
|
} |
|
|
|
go func() { |
|
|
|
log.Printf("Starting App HTTP server on %s", config.Cfg.Server.AppListenPort) |
|
|
|
if err := httpServer.ListenAndServe(); err != nil && err != http.ErrServerClosed { |
|
|
|
log.Fatalf("App server ListenAndServe error: %v", err) |
|
|
|
} |
|
|
|
}() |
|
|
|
|
|
|
|
// 5. 设置并等待优雅停机
|
|
|
|
shutdownChan := make(chan os.Signal, 1) |
|
|
|
signal.Notify(shutdownChan, syscall.SIGINT, syscall.SIGTERM) |
|
|
|
sig := <-shutdownChan |
|
|
|
log.Printf("Shutdown signal received (%s), starting graceful shutdown...", sig) |
|
|
|
|
|
|
|
// 6. 执行清理操作
|
|
|
|
shutdownCtx, cancel := context.WithTimeout(context.Background(), 10*time.Second) |
|
|
|
defer cancel() |
|
|
|
|
|
|
|
// a. 向调度服务(Redis)注销自己
|
|
|
|
registry.Unregister() |
|
|
|
|
|
|
|
// b. 优雅地关闭 HTTP 服务器
|
|
|
|
if err := httpServer.Shutdown(shutdownCtx); err != nil { |
|
|
|
log.Printf("HTTP server shutdown error: %v", err) |
|
|
|
} else { |
|
|
|
log.Println("HTTP server gracefully stopped.") |
|
|
|
} |
|
|
|
|
|
|
|
log.Println("Graceful shutdown complete.") |
|
|
|
} |
|
|
|
|
|
|
|
func listenForDevices(addr string) { |
|
|
|
@ -87,9 +129,7 @@ func listenForDevices(addr string) { |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
// ==============================================================================
|
|
|
|
// 4. 设备端认证与会话管理
|
|
|
|
// ==============================================================================
|
|
|
|
func handleDeviceSession(conn net.Conn) { |
|
|
|
defer conn.Close() |
|
|
|
log.Printf("New device connected from %s, awaiting authentication...\n", conn.RemoteAddr()) |
|
|
|
@ -117,11 +157,11 @@ func handleDeviceSession(conn net.Conn) { |
|
|
|
userID := claims.UserID |
|
|
|
log.Printf("Device '%s' (user: %s) authenticated successfully.\n", deviceSN, userID) |
|
|
|
|
|
|
|
config := yamux.DefaultConfig() |
|
|
|
config.EnableKeepAlive = true |
|
|
|
config.KeepAliveInterval = 30 * time.Second |
|
|
|
yamuxConfig := yamux.DefaultConfig() |
|
|
|
yamuxConfig.EnableKeepAlive = true |
|
|
|
yamuxConfig.KeepAliveInterval = 30 * time.Second |
|
|
|
|
|
|
|
session, err := yamux.Server(conn, config) |
|
|
|
session, err := yamux.Server(conn, yamuxConfig) |
|
|
|
if err != nil { |
|
|
|
log.Printf("Failed to start yamux session for device '%s': %v", deviceSN, err) |
|
|
|
return |
|
|
|
@ -138,6 +178,15 @@ func handleDeviceSession(conn net.Conn) { |
|
|
|
sessionMutex.Unlock() |
|
|
|
log.Printf("Yamux session started for device '%s'\n", deviceSN) |
|
|
|
|
|
|
|
if storage.RedisClient != nil { |
|
|
|
instanceID := config.Cfg.Server.InstanceID |
|
|
|
if err := storage.RedisClient.HSet(context.Background(), config.Cfg.Redis.DeviceRelayMappingKey, deviceSN, instanceID).Err(); err != nil { |
|
|
|
log.Printf("ERROR: Failed to update device-relay mapping for %s: %v", deviceSN, err) |
|
|
|
} else { |
|
|
|
log.Printf("Device %s is now mapped to instance %s in Redis.", deviceSN, instanceID) |
|
|
|
} |
|
|
|
} |
|
|
|
|
|
|
|
defer func() { |
|
|
|
sessionMutex.Lock() |
|
|
|
if currentInfo, exists := deviceSessions[deviceSN]; exists && currentInfo.Session == session { |
|
|
|
@ -145,6 +194,20 @@ func handleDeviceSession(conn net.Conn) { |
|
|
|
} |
|
|
|
sessionMutex.Unlock() |
|
|
|
log.Printf("Device '%s' session closed\n", deviceSN) |
|
|
|
|
|
|
|
// b. 再清理 Redis 映射
|
|
|
|
if storage.RedisClient != nil { |
|
|
|
instanceID := config.Cfg.Server.InstanceID |
|
|
|
// [健壮性优化] 在删除前，先检查一下 Redis 里的值是不是还是自己。
|
|
|
|
// 这可以防止因为竞态条件，错误地删除了一个刚刚重连到本机的、更新的会话映射。
|
|
|
|
currentInstanceID, err := storage.RedisClient.HGet(context.Background(), config.Cfg.Redis.DeviceRelayMappingKey, deviceSN).Result() |
|
|
|
if err == nil && currentInstanceID == instanceID { |
|
|
|
storage.RedisClient.HDel(context.Background(), config.Cfg.Redis.DeviceRelayMappingKey, deviceSN) |
|
|
|
log.Printf("Removed device-relay mapping for %s.", deviceSN) |
|
|
|
} else if err != nil && err != redis.Nil { |
|
|
|
log.Printf("ERROR: Could not verify mapping for %s before deleting: %v", deviceSN, err) |
|
|
|
} |
|
|
|
} |
|
|
|
}() |
|
|
|
|
|
|
|
<-session.CloseChan() |
|
|
|
@ -172,9 +235,7 @@ func validateDeviceToken(tokenString string) (*DeviceJWTClaims, error) { |
|
|
|
return claims, nil |
|
|
|
} |
|
|
|
|
|
|
|
// ==============================================================================
|
|
|
|
// 5. App 端认证与请求处理
|
|
|
|
// ==============================================================================
|
|
|
|
func handleAppRequest(w http.ResponseWriter, r *http.Request) { |
|
|
|
pathParts := strings.SplitN(strings.TrimPrefix(r.URL.Path, "/"), "/", 3) |
|
|
|
if len(pathParts) < 2 || pathParts[0] != "tunnel" { |
|
|
|
|
lin_hl
2 weeks ago