mirror of https://github.com/containers/podman.git
				
				
				
			
		
			
				
	
	
		
			476 lines
		
	
	
		
			8.4 KiB
		
	
	
	
		
			Go
		
	
	
	
			
		
		
	
	
			476 lines
		
	
	
		
			8.4 KiB
		
	
	
	
		
			Go
		
	
	
	
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{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) (int, 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)
 | 
						|
	n, err := buf.WriteTo(w)
 | 
						|
	return int(n), err
 | 
						|
}
 | 
						|
 | 
						|
// -----------------------------------------------------------------------------
 | 
						|
 | 
						|
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()
 | 
						|
}
 |