vscode-ssh.go

// main
/**
 * A golang ssh server which supports vscode remote ssh plugin.
 *
 * ssh_config:
 *   Host upterm
 *    HostName 192.168.101.135
 *    Port 2224
 *
 * VSCode remote ssh plugin use this command to setup its server, and use dynamic port fowarding to connect to this server.
 * 	ssh -T -D 1234 192.168.101.135 -p 2224  bash
 */
package main

import (
	"fmt"
	"io"
	"log"
	"os"
	"os/exec"
	"syscall"
	"unsafe"

	"github.com/gliderlabs/ssh"
	"github.com/kr/pty"
)

func setWinsize(f *os.File, w, h int) {
	syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), uintptr(syscall.TIOCSWINSZ),
		uintptr(unsafe.Pointer(&struct{ h, w, x, y uint16 }{uint16(h), uint16(w), 0, 0})))
}

func main() {
	svr := ssh.Server{
		Addr: ":2224",
		LocalPortForwardingCallback: ssh.LocalPortForwardingCallback(func(ctx ssh.Context, dhost string, dport uint32) bool {
			log.Println("Accepted forward", dhost, dport)
			return true
		}),
		Handler: ssh.Handler(func(s ssh.Session) {
			ptyReq, winCh, isPty := s.Pty()
			if isPty {
				cmd := exec.Command("/bin/bash")

				cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", ptyReq.Term))
				f, err := pty.Start(cmd)
				if err != nil {
					panic(err)
				}
				go func() {
					for win := range winCh {
						setWinsize(f, win.Width, win.Height)
					}
				}()
				go func() {
					io.Copy(f, s) // stdin
				}()
				io.Copy(s, f) // stdout
				cmd.Wait()

			} else {
				cmds := s.Command()
				if len(cmds) == 0 {
					cmds = []string{"/bin/bash"}
				}

				cmd := exec.Command(cmds[0], cmds[1:]...)
				cmd.Stdout = s
				cmd.Stderr = s
				cmd.Stdin = s
				err := cmd.Start()
				if err != nil {
					log.Printf("could not start command (%s)", err)
				}

				_, err = cmd.Process.Wait()
				if err != nil {
					log.Printf("failed to exit bash (%s)", err)
				}
				log.Printf("session closed")
				s.Close()
			}
		}),
		ChannelHandlers: map[string]ssh.ChannelHandler{
			"session":      ssh.DefaultSessionHandler,
			"direct-tcpip": ssh.DirectTCPIPHandler,
		},
	}
	svr.SetOption(
		ssh.HostKeyFile("/home/corvo/.ssh/id_rsa_tencent_4096"),
	)
	log.Println("starting ssh server on port 2224...")
	log.Fatal(svr.ListenAndServe())
}