test_vim_session.go

package main

import (
	"encoding/base64"
	"fmt"
	"log"
	"net/url"
	"os"
	"os/signal"
	"time"

	"github.com/gorilla/websocket"
)

func main() {
	// 获取token
	token := os.Getenv("PROCMAN_ADMIN_TOKEN")
	if token == "" {
		token = "test123" // 默认token
	}

	// 测试session_key功能
	sessionKey := "vim-test-session"

	// 第一次连接：使用指定的session_key创建新会话并启动vim
	fmt.Printf("=== 第一次连接：使用session_key '%s' 创建会话并启动vim ===\n", sessionKey)
	err := connectAndStartVim(token, sessionKey)
	if err != nil {
		log.Printf("第一次连接失败: %v", err)
		return
	}

	// 等待一段时间让vim运行
	fmt.Println("\n等待5秒让vim运行...")
	time.Sleep(5 * time.Second)

	// 第二次连接：使用相同的session_key恢复会话
	fmt.Printf("=== 第二次连接：使用session_key '%s' 恢复会话 ===\n", sessionKey)
	err = connectAndRestoreSession(token, sessionKey)
	if err != nil {
		log.Printf("第二次连接失败: %v", err)
	}
}

func connectAndStartVim(token, sessionKey string) error {
	// 创建WebSocket URL
	u := url.URL{
		Scheme: "ws",
		Host:   "localhost:8080",
		Path:   "/terminal/ws",
	}

	// 添加查询参数
	q := u.Query()
	q.Set("token", token)
	q.Set("rows", "24")
	q.Set("cols", "80")
	q.Set("command", "bash")
	q.Set("session_key", sessionKey)
	u.RawQuery = q.Encode()

	fmt.Printf("连接到: %s\n", u.String())

	// 连接WebSocket
	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		return fmt.Errorf("dial error: %w", err)
	}
	defer c.Close()

	fmt.Println("连接成功！")

	// 启动读取goroutine
	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			_, message, err := c.ReadMessage()
			if err != nil {
				log.Println("读取错误:", err)
				return
			}

			// 解析消息
			if len(message) > 0 {
				msgType := message[0]
				msgData := message[1:]
				switch msgType {
				case '1': // 输出
					if decoded, err := base64.StdEncoding.DecodeString(string(msgData)); err == nil {
						fmt.Printf("输出: %s", string(decoded))
					}
				case '2': // pong
					fmt.Println("收到Pong响应")
				case '3': // 窗口标题
					fmt.Printf("窗口标题: %s\n", string(msgData))
				}
			}
		}
	}()

	// 发送命令启动vim
	commands := []string{
		"echo 'Starting vim in 3 seconds...'",
		"sleep 3",
		"vim --version | head -5",
		"echo 'Now starting vim...'",
		"vim /tmp/test.txt", // 启动vim
	}

	for _, cmd := range commands {
		// 发送命令
		cmdWithNewline := cmd + "\n"
		encodedCmd := base64.StdEncoding.EncodeToString([]byte(cmdWithNewline))
		message := "1" + encodedCmd

		err := c.WriteMessage(websocket.TextMessage, []byte(message))
		if err != nil {
			log.Println("写入错误:", err)
			return err
		}

		// 等待命令执行
		time.Sleep(1 * time.Second)
	}

	// 等待一段时间让vim启动
	fmt.Println("等待vim启动...")
	time.Sleep(3 * time.Second)

	// 发送vim退出命令（:q!）
	fmt.Println("发送vim退出命令...")
	quitCmd := ":q!\n"
	encodedQuit := base64.StdEncoding.EncodeToString([]byte(quitCmd))
	quitMessage := "1" + encodedQuit
	err = c.WriteMessage(websocket.TextMessage, []byte(quitMessage))
	if err != nil {
		log.Println("写入错误:", err)
		return err
	}

	// 等待vim退出
	time.Sleep(2 * time.Second)

	// 关闭连接
	fmt.Println("关闭连接...")
	c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
	select {
	case <-done:
	case <-time.After(time.Second):
	}

	return nil
}

func connectAndRestoreSession(token, sessionKey string) error {
	// 创建WebSocket URL
	u := url.URL{
		Scheme: "ws",
		Host:   "localhost:8080",
		Path:   "/terminal/ws",
	}

	// 添加查询参数
	q := u.Query()
	q.Set("token", token)
	q.Set("rows", "24")
	q.Set("cols", "80")
	q.Set("command", "bash")
	q.Set("session_key", sessionKey)
	u.RawQuery = q.Encode()

	fmt.Printf("连接到: %s\n", u.String())

	// 连接WebSocket
	c, _, err := websocket.DefaultDialer.Dial(u.String(), nil)
	if err != nil {
		return fmt.Errorf("dial error: %w", err)
	}
	defer c.Close()

	fmt.Println("连接成功！检查是否恢复之前的会话...")

	// 启动读取goroutine
	done := make(chan struct{})
	go func() {
		defer close(done)
		for {
			_, message, err := c.ReadMessage()
			if err != nil {
				log.Println("读取错误:", err)
				return
			}

			// 解析消息
			if len(message) > 0 {
				msgType := message[0]
				msgData := message[1:]
				switch msgType {
				case '1': // 输出
					if decoded, err := base64.StdEncoding.DecodeString(string(msgData)); err == nil {
						output := string(decoded)
						// 检查是否包含vim相关的输出
						if len(output) > 0 {
							fmt.Printf("输出: %s", output)
						}
					}
				case '2': // pong
					// 不打印pong响应
				case '3': // 窗口标题
					fmt.Printf("窗口标题: %s\n", string(msgData))
				}
			}
		}
	}()

	// 发送一些命令来测试会话状态
	testCommands := []string{
		"echo '=== Session Restored ==='",
		"pwd",
		"echo 'Checking if vim is still running...'",
		"ps aux | grep vim",
		"echo 'Current working directory contents:'",
		"ls -la",
		"echo '=== Test Complete ==='",
	}

	for _, cmd := range testCommands {
		// 发送命令
		cmdWithNewline := cmd + "\n"
		encodedCmd := base64.StdEncoding.EncodeToString([]byte(cmdWithNewline))
		message := "1" + encodedCmd

		err := c.WriteMessage(websocket.TextMessage, []byte(message))
		if err != nil {
			log.Println("写入错误:", err)
			return err
		}

		// 等待命令执行
		time.Sleep(500 * time.Millisecond)
	}

	// 等待用户中断或超时
	interrupt := make(chan os.Signal, 1)
	signal.Notify(interrupt, os.Interrupt)

	select {
	case <-interrupt:
		fmt.Println("\n收到中断信号，关闭连接...")
		c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		select {
		case <-done:
		case <-time.After(time.Second):
		}
	case <-time.After(10 * time.Second):
		fmt.Println("\n超时，关闭连接...")
		c.WriteMessage(websocket.CloseMessage, websocket.FormatCloseMessage(websocket.CloseNormalClosure, ""))
		select {
		case <-done:
		case <-time.After(time.Second):
		}
	}

	return nil
}