✨ v0.9.21 WebSSH

cmd/agent/main.go

@@ -3,20 +3,27 @@ package main
 import (
 	"context"
 	"crypto/tls"
+	"encoding/json"
 	"errors"
 	"flag"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
 	"net/url"
 	"os"
 	"os/exec"
+	"runtime"
+	"syscall"
 	"time"
+	"unsafe"
 
 	"github.com/blang/semver"
 	"github.com/genkiroid/cert"
 	"github.com/go-ping/ping"
+	"github.com/gorilla/websocket"
+	"github.com/kr/pty"
 	"github.com/p14yground/go-github-selfupdate/selfupdate"
 	"google.golang.org/grpc"
 
@@ -58,6 +65,13 @@ const (
 	networkTimeOut = time.Second * 5  // 普通网络超时
 )
 
+type windowSize struct {
+	Rows uint16 `json:"rows"`
+	Cols uint16 `json:"cols"`
+	X    uint16
+	Y    uint16
+}
+
 func main() {
 	// 来自于 GoReleaser 的版本号
 	monitor.Version = version
@@ -154,7 +168,14 @@ func receiveTasks(tasks pb.NezhaService_RequestTaskClient) error {
 		if err != nil {
 			return err
 		}
-		go doTask(task)
+		go func() {
+			defer func() {
+				if recover() != nil {
+					println("task panic", task)
+				}
+			}()
+			doTask(task)
+		}()
 	}
 }
 
@@ -163,100 +184,16 @@ func doTask(task *pb.Task) {
 	result.Id = task.GetId()
 	result.Type = task.GetType()
 	switch task.GetType() {
+	case model.TaskTypeTerminal:
+		handleTerminalTask(task)
 	case model.TaskTypeHTTPGET:
-		start := time.Now()
-		resp, err := httpClient.Get(task.GetData())
-		if err == nil {
-			// 检查 HTTP Response 状态
-			result.Delay = float32(time.Since(start).Microseconds()) / 1000.0
-			if resp.StatusCode > 399 || resp.StatusCode < 200 {
-				err = errors.New("\n应用错误：" + resp.Status)
-			}
-		}
-		if err == nil {
-			// 检查 SSL 证书信息
-			serviceUrl, err := url.Parse(task.GetData())
-			if err == nil {
-				if serviceUrl.Scheme == "https" {
-					c := cert.NewCert(serviceUrl.Host)
-					if c.Error != "" {
-						result.Data = "SSL证书错误：" + c.Error
-					} else {
-						result.Data = c.Issuer + "|" + c.NotAfter
-						result.Successful = true
-					}
-				} else {
-					result.Successful = true
-				}
-			} else {
-				result.Data = "URL解析错误：" + err.Error()
-			}
-		} else {
-			// HTTP 请求失败
-			result.Data = err.Error()
-		}
+		handleHttpGetTask(task, &result)
 	case model.TaskTypeICMPPing:
-		pinger, err := ping.NewPinger(task.GetData())
-		if err == nil {
-			pinger.SetPrivileged(true)
-			pinger.Count = 5
-			pinger.Timeout = time.Second * 20
-			err = pinger.Run() // Blocks until finished.
-		}
-		if err == nil {
-			result.Delay = float32(pinger.Statistics().AvgRtt.Microseconds()) / 1000.0
-			result.Successful = true
-		} else {
-			result.Data = err.Error()
-		}
+		handleIcmpPingTask(task, &result)
 	case model.TaskTypeTCPPing:
-		start := time.Now()
-		conn, err := net.DialTimeout("tcp", task.GetData(), time.Second*10)
-		if err == nil {
-			conn.Write([]byte("ping\n"))
-			conn.Close()
-			result.Delay = float32(time.Since(start).Microseconds()) / 1000.0
-			result.Successful = true
-		} else {
-			result.Data = err.Error()
-		}
+		handleTcpPingTask(task, &result)
 	case model.TaskTypeCommand:
-		startedAt := time.Now()
-		var cmd *exec.Cmd
-		var endCh = make(chan struct{})
-		pg, err := utils.NewProcessExitGroup()
-		if err != nil {
-			// 进程组创建失败，直接退出
-			result.Data = err.Error()
-			client.ReportTask(context.Background(), &result)
-			return
-		}
-		timeout := time.NewTimer(time.Hour * 2)
-		if utils.IsWindows() {
-			cmd = exec.Command("cmd", "/c", task.GetData())
-		} else {
-			cmd = exec.Command("sh", "-c", task.GetData())
-		}
-		pg.AddProcess(cmd)
-		go func() {
-			select {
-			case <-timeout.C:
-				result.Data = "任务执行超时\n"
-				close(endCh)
-				pg.Dispose()
-			case <-endCh:
-				timeout.Stop()
-			}
-		}()
-		output, err := cmd.Output()
-		if err != nil {
-			result.Data += fmt.Sprintf("%s\n%s", string(output), err.Error())
-		} else {
-			close(endCh)
-			result.Data = string(output)
-			result.Successful = true
-		}
-		result.Delay = float32(time.Since(startedAt).Seconds())
+		handleCommandTask(task, &result)
 	default:
 		println("Unknown action: ", task)
 	}
@@ -307,6 +244,211 @@ func doSelfUpdate() {
 	}
 }
 
+func handleTcpPingTask(task *pb.Task, result *pb.TaskResult) {
+	start := time.Now()
+	conn, err := net.DialTimeout("tcp", task.GetData(), time.Second*10)
+	if err == nil {
+		conn.Write([]byte("ping\n"))
+		conn.Close()
+		result.Delay = float32(time.Since(start).Microseconds()) / 1000.0
+		result.Successful = true
+	} else {
+		result.Data = err.Error()
+	}
+}
+
+func handleIcmpPingTask(task *pb.Task, result *pb.TaskResult) {
+	pinger, err := ping.NewPinger(task.GetData())
+	if err == nil {
+		pinger.SetPrivileged(true)
+		pinger.Count = 5
+		pinger.Timeout = time.Second * 20
+		err = pinger.Run() // Blocks until finished.
+	}
+	if err == nil {
+		result.Delay = float32(pinger.Statistics().AvgRtt.Microseconds()) / 1000.0
+		result.Successful = true
+	} else {
+		result.Data = err.Error()
+	}
+}
+
+func handleHttpGetTask(task *pb.Task, result *pb.TaskResult) {
+	start := time.Now()
+	resp, err := httpClient.Get(task.GetData())
+	if err == nil {
+		// 检查 HTTP Response 状态
+		result.Delay = float32(time.Since(start).Microseconds()) / 1000.0
+		if resp.StatusCode > 399 || resp.StatusCode < 200 {
+			err = errors.New("\n应用错误：" + resp.Status)
+		}
+	}
+	if err == nil {
+		// 检查 SSL 证书信息
+		serviceUrl, err := url.Parse(task.GetData())
+		if err == nil {
+			if serviceUrl.Scheme == "https" {
+				c := cert.NewCert(serviceUrl.Host)
+				if c.Error != "" {
+					result.Data = "SSL证书错误：" + c.Error
+				} else {
+					result.Data = c.Issuer + "|" + c.NotAfter
+					result.Successful = true
+				}
+			} else {
+				result.Successful = true
+			}
+		} else {
+			result.Data = "URL解析错误：" + err.Error()
+		}
+	} else {
+		// HTTP 请求失败
+		result.Data = err.Error()
+	}
+}
+
+func handleCommandTask(task *pb.Task, result *pb.TaskResult) {
+	startedAt := time.Now()
+	var cmd *exec.Cmd
+	var endCh = make(chan struct{})
+	pg, err := utils.NewProcessExitGroup()
+	if err != nil {
+		// 进程组创建失败，直接退出
+		result.Data = err.Error()
+		return
+	}
+	timeout := time.NewTimer(time.Hour * 2)
+	if utils.IsWindows() {
+		cmd = exec.Command("cmd", "/c", task.GetData())
+	} else {
+		cmd = exec.Command("sh", "-c", task.GetData())
+	}
+	pg.AddProcess(cmd)
+	go func() {
+		select {
+		case <-timeout.C:
+			result.Data = "任务执行超时\n"
+			close(endCh)
+			pg.Dispose()
+		case <-endCh:
+			timeout.Stop()
+		}
+	}()
+	output, err := cmd.Output()
+	if err != nil {
+		result.Data += fmt.Sprintf("%s\n%s", string(output), err.Error())
+	} else {
+		close(endCh)
+		result.Data = string(output)
+		result.Successful = true
+	}
+	result.Delay = float32(time.Since(startedAt).Seconds())
+}
+
+func handleTerminalTask(task *pb.Task) {
+	var terminal model.TerminalTask
+	err := json.Unmarshal([]byte(task.GetData()), &terminal)
+	if err != nil {
+		println("Terminal 任务解析错误：", err)
+		return
+	}
+	protocol := "ws"
+	if terminal.UseSSL {
+		protocol += "s"
+	}
+	header := http.Header{}
+	header.Add("Secret", clientSecret)
+	conn, _, err := websocket.DefaultDialer.Dial(fmt.Sprintf("%s://%s/terminal/%s", protocol, terminal.Host, terminal.Session), header)
+	if err != nil {
+		println("Terminal 连接失败：", err)
+		return
+	}
+	defer conn.Close()
+
+	var cmd *exec.Cmd
+	var shellPath string
+	if runtime.GOOS == "windows" {
+		shellPath, err = exec.LookPath("powershell.exe")
+		if err != nil || shellPath == "" {
+			shellPath = "cmd.exe"
+		}
+	} else {
+		shellPath = os.Getenv("SHELL")
+		if shellPath == "" {
+			shellPath = "sh"
+		}
+	}
+	cmd = exec.Command(shellPath)
+	cmd.Env = append(os.Environ(), "TERM=xterm")
+
+	tty, err := pty.Start(cmd)
+	if err != nil {
+		println("Terminal pty.Start失败：", err)
+		return
+	}
+
+	defer func() {
+		cmd.Process.Kill()
+		cmd.Process.Wait()
+		tty.Close()
+		conn.Close()
+		println("terminal exit", terminal.Session)
+	}()
+	println("terminal init", terminal.Session, shellPath)
+
+	go func() {
+		for {
+			buf := make([]byte, 1024)
+			read, err := tty.Read(buf)
+			if err != nil {
+				conn.WriteMessage(websocket.TextMessage, []byte(err.Error()))
+				return
+			}
+			conn.WriteMessage(websocket.BinaryMessage, buf[:read])
+		}
+	}()
+
+	for {
+		messageType, reader, err := conn.NextReader()
+		if err != nil {
+			return
+		}
+
+		if messageType == websocket.TextMessage {
+			continue
+		}
+
+		dataTypeBuf := make([]byte, 1)
+		read, err := reader.Read(dataTypeBuf)
+		if err != nil {
+			conn.WriteMessage(websocket.TextMessage, []byte("Unable to read message type from reader"))
+			return
+		}
+
+		if read != 1 {
+			return
+		}
+
+		switch dataTypeBuf[0] {
+		case 0:
+			io.Copy(tty, reader)
+		case 1:
+			decoder := json.NewDecoder(reader)
+			resizeMessage := windowSize{}
+			err := decoder.Decode(&resizeMessage)
+			if err != nil {
+				continue
+			}
+			syscall.Syscall(
+				syscall.SYS_IOCTL,
+				tty.Fd(),
+				syscall.TIOCSWINSZ,
+				uintptr(unsafe.Pointer(&resizeMessage)),
+			)
+		}
+	}
+}
+
 func println(v ...interface{}) {
 	if debug {
 		log.Println(v...)

cmd/dashboard/controller/common_page.go

@@ -1,22 +1,37 @@
 package controller
 
 import (
+	"encoding/json"
 	"errors"
 	"fmt"
+	"log"
 	"net/http"
+	"sync"
 	"time"
 
 	"github.com/gin-gonic/gin"
 	"github.com/gorilla/websocket"
+	"github.com/hashicorp/go-uuid"
 	"golang.org/x/crypto/bcrypt"
 
 	"github.com/naiba/nezha/model"
 	"github.com/naiba/nezha/pkg/mygin"
+	"github.com/naiba/nezha/proto"
 	"github.com/naiba/nezha/service/dao"
 )
 
+type terminalContext struct {
+	agentConn *websocket.Conn
+	userConn  *websocket.Conn
+	serverID  uint64
+	host      string
+	useSSL    bool
+}
+
 type commonPage struct {
-	r *gin.Engine
+	r             *gin.Engine
+	terminals     map[string]*terminalContext
+	terminalsLock *sync.Mutex
 }
 
 func (cp *commonPage) serve() {
@@ -27,6 +42,8 @@ func (cp *commonPage) serve() {
 	cr.GET("/", cp.home)
 	cr.GET("/service", cp.service)
 	cr.GET("/ws", cp.ws)
+	cr.POST("/terminal", cp.createTerminal)
+	cr.GET("/terminal/:id", cp.terminal)
 }
 
 type viewPasswordForm struct {
@@ -98,7 +115,10 @@ func (cp *commonPage) home(c *gin.Context) {
 	}))
 }
 
-var upgrader = websocket.Upgrader{}
+var upgrader = websocket.Upgrader{
+	ReadBufferSize:  1024,
+	WriteBufferSize: 1024,
+}
 
 type Data struct {
 	Now     int64           `json:"now,omitempty"`
@@ -139,3 +159,217 @@ func (cp *commonPage) ws(c *gin.Context) {
 		time.Sleep(time.Second * 2)
 	}
 }
+
+func (cp *commonPage) terminal(c *gin.Context) {
+	log.Println("terminal connected", c.Request.URL)
+	defer log.Println("terminal disconnected", c.Request.URL)
+	terminalID := c.Param("id")
+	cp.terminalsLock.Lock()
+	if terminalID == "" || cp.terminals[terminalID] == nil {
+		cp.terminalsLock.Unlock()
+		mygin.ShowErrorPage(c, mygin.ErrInfo{
+			Code:  http.StatusForbidden,
+			Title: "无权访问",
+			Msg:   "终端会话不存在",
+			Link:  "/",
+			Btn:   "返回首页",
+		}, true)
+		return
+	}
+
+	terminal := cp.terminals[terminalID]
+	cp.terminalsLock.Unlock()
+
+	defer func() {
+		cp.terminalsLock.Lock()
+		defer cp.terminalsLock.Unlock()
+		delete(cp.terminals, terminalID)
+	}()
+
+	var isAgent bool
+
+	if _, authorized := c.Get(model.CtxKeyAuthorizedUser); !authorized {
+		dao.ServerLock.RLock()
+		_, hasID := dao.SecretToID[c.Request.Header.Get("Secret")]
+		dao.ServerLock.RUnlock()
+		if !hasID {
+			mygin.ShowErrorPage(c, mygin.ErrInfo{
+				Code:  http.StatusForbidden,
+				Title: "无权访问",
+				Msg:   "用户未登录或非法终端",
+				Link:  "/",
+				Btn:   "返回首页",
+			}, true)
+			return
+		}
+		if terminal.userConn == nil {
+			mygin.ShowErrorPage(c, mygin.ErrInfo{
+				Code:  http.StatusForbidden,
+				Title: "无权访问",
+				Msg:   "用户不在线",
+				Link:  "/",
+				Btn:   "返回首页",
+			}, true)
+			return
+		}
+		if terminal.agentConn != nil {
+			mygin.ShowErrorPage(c, mygin.ErrInfo{
+				Code:  http.StatusInternalServerError,
+				Title: "连接已存在",
+				Msg:   "Websocket协议切换失败",
+				Link:  "/",
+				Btn:   "返回首页",
+			}, true)
+			return
+		}
+		isAgent = true
+	} else {
+		dao.ServerLock.RLock()
+		server := dao.ServerList[terminal.serverID]
+		dao.ServerLock.RUnlock()
+		if server == nil || server.TaskStream == nil {
+			mygin.ShowErrorPage(c, mygin.ErrInfo{
+				Code:  http.StatusForbidden,
+				Title: "请求失败",
+				Msg:   "服务器不存在或处于离线状态",
+				Link:  "/server",
+				Btn:   "返回重试",
+			}, true)
+			return
+		}
+
+		terminalData, _ := json.Marshal(&model.TerminalTask{
+			Host:    terminal.host,
+			UseSSL:  terminal.useSSL,
+			Session: terminalID,
+		})
+
+		if err := server.TaskStream.Send(&proto.Task{
+			Type: model.TaskTypeTerminal,
+			Data: string(terminalData),
+		}); err != nil {
+			mygin.ShowErrorPage(c, mygin.ErrInfo{
+				Code:  http.StatusForbidden,
+				Title: "请求失败",
+				Msg:   "Agent信令下发失败",
+				Link:  "/server",
+				Btn:   "返回重试",
+			}, true)
+			return
+		}
+	}
+
+	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		mygin.ShowErrorPage(c, mygin.ErrInfo{
+			Code:  http.StatusInternalServerError,
+			Title: "网络错误",
+			Msg:   "Websocket协议切换失败",
+			Link:  "/",
+			Btn:   "返回首页",
+		}, true)
+		return
+	}
+	defer conn.Close()
+
+	if isAgent {
+		terminal.agentConn = conn
+	} else {
+		terminal.userConn = conn
+		defer func() {
+			// 用户断开链接时断开 Agent 连接
+			if terminal.agentConn != nil {
+				terminal.agentConn.Close()
+			}
+		}()
+	}
+
+	for {
+		msgType, data, err := conn.ReadMessage()
+		if err != nil {
+			return
+		}
+		// 将文本消息转换为命令输入
+		if msgType == websocket.TextMessage {
+			data = append([]byte{0}, data...)
+		}
+		// 传递给对方
+		if isAgent {
+			err = terminal.userConn.WriteMessage(websocket.BinaryMessage, data)
+		} else {
+			err = terminal.agentConn.WriteMessage(websocket.BinaryMessage, data)
+		}
+		if err != nil {
+			return
+		}
+	}
+}
+
+type createTerminalRequest struct {
+	Host     string
+	Protocol string
+	ID       uint64
+}
+
+func (cp *commonPage) createTerminal(c *gin.Context) {
+	if _, authorized := c.Get(model.CtxKeyAuthorizedUser); !authorized {
+		mygin.ShowErrorPage(c, mygin.ErrInfo{
+			Code:  http.StatusForbidden,
+			Title: "无权访问",
+			Msg:   "用户未登录",
+			Link:  "/login",
+			Btn:   "去登录",
+		}, true)
+		return
+	}
+	var createTerminalReq createTerminalRequest
+	if err := c.ShouldBind(&createTerminalReq); err != nil {
+		mygin.ShowErrorPage(c, mygin.ErrInfo{
+			Code:  http.StatusForbidden,
+			Title: "请求失败",
+			Msg:   "请求参数有误：" + err.Error(),
+			Link:  "/server",
+			Btn:   "返回重试",
+		}, true)
+		return
+	}
+
+	id, err := uuid.GenerateUUID()
+	if err != nil {
+		mygin.ShowErrorPage(c, mygin.ErrInfo{
+			Code:  http.StatusInternalServerError,
+			Title: "系统错误",
+			Msg:   "生成会话ID失败",
+			Link:  "/server",
+			Btn:   "返回重试",
+		}, true)
+		return
+	}
+
+	dao.ServerLock.RLock()
+	server := dao.ServerList[createTerminalReq.ID]
+	dao.ServerLock.RUnlock()
+	if server == nil || server.TaskStream == nil {
+		mygin.ShowErrorPage(c, mygin.ErrInfo{
+			Code:  http.StatusForbidden,
+			Title: "请求失败",
+			Msg:   "服务器不存在或处于离线状态",
+			Link:  "/server",
+			Btn:   "返回重试",
+		}, true)
+		return
+	}
+
+	cp.terminalsLock.Lock()
+	defer cp.terminalsLock.Unlock()
+
+	cp.terminals[id] = &terminalContext{
+		serverID: createTerminalReq.ID,
+		host:     createTerminalReq.Host,
+		useSSL:   createTerminalReq.Protocol == "https:",
+	}
+
+	c.HTML(http.StatusOK, "dashboard/terminal", mygin.CommonEnvironment(c, gin.H{
+		"SessionID": id,
+	}))
+}

cmd/dashboard/controller/controller.go

@@ -5,6 +5,7 @@ import (
 	"html/template"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
 	"code.cloudfoundry.org/bytefmt"
@@ -135,7 +136,7 @@ func ServeWeb(port uint) *http.Server {
 
 func routers(r *gin.Engine) {
 	// 通用页面
-	cp := commonPage{r}
+	cp := commonPage{r: r, terminals: make(map[string]*terminalContext), terminalsLock: new(sync.Mutex)}
 	cp.serve()
 	// 游客页面
 	gp := guestPage{r}

cmd/playground/main.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"crypto/tls"
 	"fmt"
+	"io"
 	"log"
 	"net"
 	"net/http"
@@ -27,8 +28,41 @@ func main() {
 	// httpWithSSLInfo()
 	// sysinfo()
 	// cmdExec()
-	resolveIP("ipapi.co", true)
-	resolveIP("ipapi.co", false)
+	// resolveIP("ipapi.co", true)
+	// resolveIP("ipapi.co", false)
+	log.Println(exec.LookPath("powershell.exe"))
+	defaultShell := os.Getenv("SHELL")
+	if defaultShell == "" {
+		defaultShell = "sh"
+	}
+	cmd := exec.Command(defaultShell)
+	cmd.Stdin = os.Stdin
+	stdoutReader, err := cmd.StdoutPipe()
+	if err != nil {
+		println("Terminal StdoutPipe:", err)
+		return
+	}
+	stderrReader, err := cmd.StderrPipe()
+	if err != nil {
+		println("Terminal StderrPipe: ", err)
+		return
+	}
+
+	readers := []io.Reader{stdoutReader, stderrReader}
+	for i := 0; i < len(readers); i++ {
+		go func(j int) {
+			data := make([]byte, 2048)
+			for {
+				count, err := readers[j].Read(data)
+				if err != nil {
+					panic(err)
+				}
+				os.Stdout.Write(data[:count])
+			}
+		}(i)
+	}
+
+	cmd.Run()
 }
 
 func resolveIP(addr string, ipv6 bool) {