adbssh.go

package main

import (
	"context"
	"crypto/rand"
	"crypto/rsa"
	"encoding/binary"
	"io"
	"log/slog"
	"net"
	"os"
	"strconv"
	"sync"

	"github.com/lmittmann/tint"
	"github.com/pgaskin/go-adb/adb/adbhost"
	"github.com/pgaskin/go-adb/adblib/adbexec/v2"
	"github.com/pgaskin/go-adb/adblib/adbnet"
	"golang.org/x/crypto/ssh"
)

/*
Host adbssh
    HostName 127.0.0.1
    Port 50372
    UserKnownHostsFile /dev/null
    StrictHostKeyChecking no
    IdentityFile ~/.ssh/id_ed25519
*/

func main() {
	slog.SetDefault(slog.New(tint.NewHandler(os.Stdout, &tint.Options{
		Level: slog.LevelDebug,
	})))

	ctx := context.Background()

	sscfg := &ssh.ServerConfig{
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			return &ssh.Permissions{}, nil // TODO: auth with existing adbkey
		},
	}
	sscfg.AddHostKey(must1(ssh.NewSignerFromKey(must1(rsa.GenerateKey(rand.Reader, 2048)))))

	dlr := &adbhost.Dialer{}
	must(dlr.LoadFeatures(ctx))

	srv := adbhost.Server(dlr, adbhost.Serial(os.Getenv("ANDROID_SERIAL"))) // TODO: use the ssh username
	must(srv.LoadFeatures(ctx))

	listener := must1(net.Listen("tcp", ":50372"))
	defer listener.Close()

	for {
		netconn := must1(listener.Accept())

		go func() {
			conn, chans, requests, err := ssh.NewServerConn(netconn, sscfg)
			if err != nil {
				slog.Error("failed to handle ssh connection", "err", err)
			}
			defer conn.Close()

			var wg sync.WaitGroup
			defer wg.Wait()

			wg.Add(1)
			go func() {
				defer wg.Done()
				for req := range requests {
					switch req.Type {
					// case "tcpip-forward": // TODO: port reverse
					// case "cancel-tcpip-forward": // TODO: port reverse
					default:
						slog.Debug("unhandled request", "type", req.Type)
						req.Reply(false, []byte("not supported"))
					}
				}
			}()

			for req := range chans {
				switch req.ChannelType() {
				case "session":
					channel, requests, err := req.Accept()
					if err != nil {
						slog.Error("failed to accept channel", "err", err)
						continue
					}

					cmd := &adbexec.Cmd{
						Server: srv,
						Stdin:  channel,
						Stdout: channel,
						Stderr: channel,
					}
					var setupPTY func()

					run := func() {
						wg.Add(1)
						go func() {
							defer wg.Done()
							defer channel.Close()

							if setupPTY != nil {
								setupPTY()
							}

							if err := cmd.Wait(); err != nil {
								code := 128
								if xx := err.(*adbexec.ExitError); xx != nil {
									if n := xx.ExitCode(); n > 0 {
										code = n
									}
								}
								var data struct {
									Status uint32
								}
								data.Status = uint32(code)
								if _, err := channel.SendRequest("exit-status", false, ssh.Marshal(&data)); err != nil {
									slog.Warn("failed to send exit status", "err", err)
								}
							}
						}()
					}

					wg.Add(1)
					go func() {
						defer wg.Done()
						for req := range requests {
							switch req.Type {
							case "env":
								req.Reply(false, []byte("not supported")) // ignore, not supported

							case "signal":
								var data struct {
									Signal string
								}
								if err := ssh.Unmarshal(req.Payload, &data); err != nil {
									req.Reply(false, []byte("failed to parse request: "+err.Error()))
									continue
								}
								slog.Warn("ignoring signal", "signal", data.Signal)
								req.Reply(false, []byte("not supported")) // ignore, not supported

							case "pty-req":
								var data struct {
									Term   string
									Col    uint32
									Row    uint32
									XPixel uint32
									YPixel uint32
								}
								if len(req.Payload) >= 4 {
									n := int(binary.BigEndian.Uint32(req.Payload[0:4]))
									if len(req.Payload) < 4+n {
										req.Reply(false, []byte("failed to parse request"))
										return
									}
									data.Term = string(req.Payload[4 : 4+n])
									if len(req.Payload) >= 4+n+8 {
										data.Col = binary.BigEndian.Uint32(req.Payload[4+n+0:])
										data.Row = binary.BigEndian.Uint32(req.Payload[4+n+4:])
									}
									if len(req.Payload) >= 4+n+16 {
										data.XPixel = binary.BigEndian.Uint32(req.Payload[4+n+8:])
										data.YPixel = binary.BigEndian.Uint32(req.Payload[4+n+12:])
									}
								}
								if cmd == nil {
									req.Reply(false, []byte("subsystem already started"))
									continue
								}
								if cmd.Process != nil {
									req.Reply(false, []byte("process already started"))
									continue
								}

								cmd.PTY = true
								cmd.Term = data.Term
								setupPTY = func() { cmd.Process.Resize(int(data.Row), int(data.Col), int(data.XPixel), int(data.YPixel)) }
								req.Reply(true, nil)

							case "window-change":
								var data struct {
									Term   string
									Col    uint32
									Row    uint32
									XPixel uint32
									YPixel uint32
								}
								if len(req.Payload) >= 8 {
									data.Col = binary.BigEndian.Uint32(req.Payload[0:4])
									data.Row = binary.BigEndian.Uint32(req.Payload[4:8])
								}
								if len(req.Payload) >= 16 {
									data.XPixel = binary.BigEndian.Uint32(req.Payload[8:12])
									data.YPixel = binary.BigEndian.Uint32(req.Payload[12:16])
								}
								if cmd == nil {
									req.Reply(false, []byte("subsystem started"))
									continue
								}
								if cmd.Process == nil {
									req.Reply(false, []byte("process not started"))
									continue
								}
								if !cmd.PTY {
									req.Reply(false, nil)
									continue
								}
								if err := cmd.Process.Resize(int(data.Row), int(data.Col), int(data.XPixel), int(data.YPixel)); err != nil {
									req.Reply(false, []byte(err.Error()))
									continue
								}
								req.Reply(true, nil)

							case "shell":
								if cmd == nil {
									req.Reply(false, []byte("subsystem already started"))
									continue
								}
								if cmd.Process != nil {
									req.Reply(false, []byte("process already started"))
									continue
								}
								if err := cmd.Start(); err != nil {
									req.Reply(false, []byte("failed to start command: "+err.Error()))
									continue
								}
								req.Reply(true, nil)

								run()

							case "exec":
								var data struct {
									Command string
								}
								if err := ssh.Unmarshal(req.Payload, &data); err != nil {
									req.Reply(false, []byte("failed to parse request: "+err.Error()))
									continue
								}
								if cmd == nil {
									req.Reply(false, []byte("subsystem already started"))
									continue
								}
								if cmd.Process != nil {
									req.Reply(false, []byte("process already started"))
									continue
								}
								cmd.Command = data.Command
								if err := cmd.Start(); err != nil {
									req.Reply(false, []byte("failed to start command: "+err.Error()))
									continue
								}
								req.Reply(true, nil)

								run()

							case "subsystem":
								var data struct {
									Name string
								}
								if err := ssh.Unmarshal(req.Payload, &data); err != nil {
									req.Reply(false, []byte("failed to parse request: "+err.Error()))
									continue
								}
								if cmd == nil {
									req.Reply(false, []byte("subsystem already started"))
									continue
								}
								cmd = nil

								switch data.Name {
								case "sftp":
									// TODO: option to push and start real sftp server on device?
									slog.Debug("rejected sftp request")
									req.Reply(false, []byte("sftp not available")) // we can't really implement sftp over adb syncproto, which is more similar to scp anyways (no random i/o, etc)
									continue

								default:
									slog.Debug("unhandled subsystem request", "name", data.Name)
									req.Reply(false, []byte("unknown subsystem"))
									continue
								}

							default:
								slog.Debug("unhandled session request", "type", req.Type)
								if req.WantReply {
									req.Reply(false, nil)
								}
							}
						}
					}()

				case "direct-tcpip":
					var data struct {
						DestAddr string
						DestPort uint32

						OriginAddr string
						OriginPort uint32
					}
					if err := ssh.Unmarshal(req.ExtraData(), &data); err != nil {
						req.Reject(ssh.ConnectionFailed, "failed to parse direct-tcpip request: "+err.Error())
						return
					}

					nconn, err := adbnet.Dial(srv, "tcp", net.JoinHostPort(data.DestAddr, strconv.FormatUint(uint64(data.DestPort), 10)))
					if err != nil {
						req.Reject(ssh.ConnectionFailed, err.Error())
						return
					}

					ch, reqs, err := req.Accept()
					if err != nil {
						nconn.Close()
						slog.Error("failed to accept channel", "err", err)
						continue
					}
					go ssh.DiscardRequests(reqs)

					go func() {
						defer ch.Close()
						defer nconn.Close()
						io.Copy(ch, nconn)
					}()

					go func() {
						defer ch.Close()
						defer nconn.Close()
						io.Copy(nconn, ch)
					}()

				default:
					slog.Debug("unhandled channel", "type", req.ChannelType())
					req.Reject(ssh.UnknownChannelType, "unknown channel type")
				}
			}
		}()
	}
}

func must(err error) {
	if err != nil {
		panic(err)
	}
}

func must1[T any](r1 T, err error) T {
	if err != nil {
		panic(err)
	}
	return r1
}