package readline import ( "bufio" "bytes" "encoding/binary" "fmt" "io" "net" "os" "sync" "sync/atomic" ) type MsgType int16 const ( T_DATA = MsgType(iota) T_WIDTH T_WIDTH_REPORT T_ISTTY_REPORT T_RAW T_ERAW // exit raw T_EOF ) type RemoteSvr struct { eof int32 closed int32 width int32 reciveChan chan struct{} writeChan chan *writeCtx conn net.Conn isTerminal bool funcWidthChan func() stopChan chan struct{} dataBufM sync.Mutex dataBuf bytes.Buffer } type writeReply struct { n int err error } type writeCtx struct { msg *Message reply chan *writeReply } func newWriteCtx(msg *Message) *writeCtx { return &writeCtx{ msg: msg, reply: make(chan *writeReply), } } func NewRemoteSvr(conn net.Conn) (*RemoteSvr, error) { rs := &RemoteSvr{ width: -1, conn: conn, writeChan: make(chan *writeCtx), reciveChan: make(chan struct{}), stopChan: make(chan struct{}), } buf := bufio.NewReader(rs.conn) if err := rs.init(buf); err != nil { return nil, err } go rs.readLoop(buf) go rs.writeLoop() return rs, nil } func (r *RemoteSvr) init(buf *bufio.Reader) error { m, err := ReadMessage(buf) if err != nil { return err } // receive isTerminal if m.Type != T_ISTTY_REPORT { return fmt.Errorf("unexpected init message") } r.GotIsTerminal(m.Data) // receive width m, err = ReadMessage(buf) if err != nil { return err } if m.Type != T_WIDTH_REPORT { return fmt.Errorf("unexpected init message") } r.GotReportWidth(m.Data) return nil } func (r *RemoteSvr) HandleConfig(cfg *Config) { cfg.Stderr = r cfg.Stdout = r cfg.Stdin = r cfg.FuncExitRaw = r.ExitRawMode cfg.FuncIsTerminal = r.IsTerminal cfg.FuncMakeRaw = r.EnterRawMode cfg.FuncExitRaw = r.ExitRawMode cfg.FuncGetWidth = r.GetWidth cfg.FuncOnWidthChanged = func(f func()) { r.funcWidthChan = f } } func (r *RemoteSvr) IsTerminal() bool { return r.isTerminal } func (r *RemoteSvr) checkEOF() error { if atomic.LoadInt32(&r.eof) == 1 { return io.EOF } return nil } func (r *RemoteSvr) Read(b []byte) (int, error) { r.dataBufM.Lock() n, err := r.dataBuf.Read(b) r.dataBufM.Unlock() if n == 0 { if err := r.checkEOF(); err != nil { return 0, err } } if n == 0 && err == io.EOF { <-r.reciveChan r.dataBufM.Lock() n, err = r.dataBuf.Read(b) r.dataBufM.Unlock() } if n == 0 { if err := r.checkEOF(); err != nil { return 0, err } } return n, err } func (r *RemoteSvr) writeMsg(m *Message) error { ctx := newWriteCtx(m) r.writeChan <- ctx reply := <-ctx.reply return reply.err } func (r *RemoteSvr) Write(b []byte) (int, error) { ctx := newWriteCtx(NewMessage(T_DATA, b)) r.writeChan <- ctx reply := <-ctx.reply return reply.n, reply.err } func (r *RemoteSvr) EnterRawMode() error { return r.writeMsg(NewMessage(T_RAW, nil)) } func (r *RemoteSvr) ExitRawMode() error { return r.writeMsg(NewMessage(T_ERAW, nil)) } func (r *RemoteSvr) writeLoop() { defer r.Close() loop: for { select { case ctx, ok := <-r.writeChan: if !ok { break } n, err := ctx.msg.WriteTo(r.conn) ctx.reply <- &writeReply{int(n), err} case <-r.stopChan: break loop } } } func (r *RemoteSvr) Close() error { if atomic.CompareAndSwapInt32(&r.closed, 0, 1) { close(r.stopChan) r.conn.Close() } return nil } func (r *RemoteSvr) readLoop(buf *bufio.Reader) { defer r.Close() for { m, err := ReadMessage(buf) if err != nil { break } switch m.Type { case T_EOF: atomic.StoreInt32(&r.eof, 1) select { case r.reciveChan <- struct{}{}: default: } case T_DATA: r.dataBufM.Lock() r.dataBuf.Write(m.Data) r.dataBufM.Unlock() select { case r.reciveChan <- struct{}{}: default: } case T_WIDTH_REPORT: r.GotReportWidth(m.Data) case T_ISTTY_REPORT: r.GotIsTerminal(m.Data) } } } func (r *RemoteSvr) GotIsTerminal(data []byte) { if binary.BigEndian.Uint16(data) == 0 { r.isTerminal = false } else { r.isTerminal = true } } func (r *RemoteSvr) GotReportWidth(data []byte) { atomic.StoreInt32(&r.width, int32(binary.BigEndian.Uint16(data))) if r.funcWidthChan != nil { r.funcWidthChan() } } func (r *RemoteSvr) GetWidth() int { return int(atomic.LoadInt32(&r.width)) } // ----------------------------------------------------------------------------- type Message struct { Type MsgType Data []byte } func ReadMessage(r io.Reader) (*Message, error) { m := new(Message) var length int32 if err := binary.Read(r, binary.BigEndian, &length); err != nil { return nil, err } if err := binary.Read(r, binary.BigEndian, &m.Type); err != nil { return nil, err } m.Data = make([]byte, int(length)-2) if _, err := io.ReadFull(r, m.Data); err != nil { return nil, err } return m, nil } func NewMessage(t MsgType, data []byte) *Message { return &Message{t, data} } func (m *Message) WriteTo(w io.Writer) (int64, error) { buf := bytes.NewBuffer(make([]byte, 0, len(m.Data)+2+4)) binary.Write(buf, binary.BigEndian, int32(len(m.Data)+2)) binary.Write(buf, binary.BigEndian, m.Type) buf.Write(m.Data) return buf.WriteTo(w) } // ----------------------------------------------------------------------------- type RemoteCli struct { conn net.Conn raw RawMode receiveChan chan struct{} inited int32 isTerminal *bool data bytes.Buffer dataM sync.Mutex } func NewRemoteCli(conn net.Conn) (*RemoteCli, error) { r := &RemoteCli{ conn: conn, receiveChan: make(chan struct{}), } return r, nil } func (r *RemoteCli) MarkIsTerminal(is bool) { r.isTerminal = &is } func (r *RemoteCli) init() error { if !atomic.CompareAndSwapInt32(&r.inited, 0, 1) { return nil } if err := r.reportIsTerminal(); err != nil { return err } if err := r.reportWidth(); err != nil { return err } // register sig for width changed DefaultOnWidthChanged(func() { r.reportWidth() }) return nil } func (r *RemoteCli) writeMsg(m *Message) error { r.dataM.Lock() _, err := m.WriteTo(r.conn) r.dataM.Unlock() return err } func (r *RemoteCli) Write(b []byte) (int, error) { m := NewMessage(T_DATA, b) r.dataM.Lock() _, err := m.WriteTo(r.conn) r.dataM.Unlock() return len(b), err } func (r *RemoteCli) reportWidth() error { screenWidth := GetScreenWidth() data := make([]byte, 2) binary.BigEndian.PutUint16(data, uint16(screenWidth)) msg := NewMessage(T_WIDTH_REPORT, data) if err := r.writeMsg(msg); err != nil { return err } return nil } func (r *RemoteCli) reportIsTerminal() error { var isTerminal bool if r.isTerminal != nil { isTerminal = *r.isTerminal } else { isTerminal = DefaultIsTerminal() } data := make([]byte, 2) if isTerminal { binary.BigEndian.PutUint16(data, 1) } else { binary.BigEndian.PutUint16(data, 0) } msg := NewMessage(T_ISTTY_REPORT, data) if err := r.writeMsg(msg); err != nil { return err } return nil } func (r *RemoteCli) readLoop() { buf := bufio.NewReader(r.conn) for { msg, err := ReadMessage(buf) if err != nil { break } switch msg.Type { case T_ERAW: r.raw.Exit() case T_RAW: r.raw.Enter() case T_DATA: os.Stdout.Write(msg.Data) } } } func (r *RemoteCli) ServeBy(source io.Reader) error { if err := r.init(); err != nil { return err } go func() { defer r.Close() for { n, _ := io.Copy(r, source) if n == 0 { break } } }() defer r.raw.Exit() r.readLoop() return nil } func (r *RemoteCli) Close() { r.writeMsg(NewMessage(T_EOF, nil)) } func (r *RemoteCli) Serve() error { return r.ServeBy(os.Stdin) } func ListenRemote(n, addr string, cfg *Config, h func(*Instance), onListen ...func(net.Listener) error) error { ln, err := net.Listen(n, addr) if err != nil { return err } if len(onListen) > 0 { if err := onListen[0](ln); err != nil { return err } } for { conn, err := ln.Accept() if err != nil { break } go func() { defer conn.Close() rl, err := HandleConn(*cfg, conn) if err != nil { return } h(rl) }() } return nil } func HandleConn(cfg Config, conn net.Conn) (*Instance, error) { r, err := NewRemoteSvr(conn) if err != nil { return nil, err } r.HandleConfig(&cfg) rl, err := NewEx(&cfg) if err != nil { return nil, err } return rl, nil } func DialRemote(n, addr string) error { conn, err := net.Dial(n, addr) if err != nil { return err } defer conn.Close() cli, err := NewRemoteCli(conn) if err != nil { return err } return cli.Serve() }