ssh: export a transport interface

Author: y3llowcake

ssh/channel.go

@@ -212,7 +212,7 @@ func (ch *channel) writePacket(packet []byte) error {
 		return io.EOF
 	}
 	ch.sentClose = (packet[0] == msgChannelClose)
-	err := ch.mux.conn.writePacket(packet)
+	err := ch.mux.conn.WritePacket(packet)
 	ch.writeMu.Unlock()
 	return err
 }

ssh/client.go

@@ -88,6 +88,13 @@ func NewClientConn(c net.Conn, addr string, config *ClientConfig) (Conn, <-chan
 	return conn, conn.mux.incomingChannels, conn.mux.incomingRequests, nil
 }
 
+func NewClientConnFromTransport(t Transport) (Conn, <-chan NewChannel, <-chan *Request, error) {
+	conn := &connection{
+		mux: newMux(t),
+	}
+	return conn, conn.mux.incomingChannels, conn.mux.incomingRequests, nil
+}
+
 // clientHandshake performs the client side key exchange. See RFC 4253 Section
 // 7.
 func (c *connection) clientHandshake(dialAddress string, config *ClientConfig) error {

ssh/handshake.go

@@ -183,6 +183,10 @@ func (t *handshakeTransport) printPacket(p []byte, write bool) {
 	}
 }
 
+func (t *handshakeTransport) ReadPacket() ([]byte, error) {
+	return t.readPacket()
+}
+
 func (t *handshakeTransport) readPacket() ([]byte, error) {
 	p, ok := <-t.incoming
 	if !ok {
@@ -479,6 +483,10 @@ func (t *handshakeTransport) sendKexInit() error {
 	return nil
 }
 
+func (t *handshakeTransport) WritePacket(p []byte) error {
+	return t.writePacket(p)
+}
+
 func (t *handshakeTransport) writePacket(p []byte) error {
 	switch p[0] {
 	case msgKexInit:

ssh/mux.go

@@ -86,7 +86,7 @@ func (c *chanList) dropAll() []*channel {
 // mux represents the state for the SSH connection protocol, which
 // multiplexes many channels onto a single packet transport.
 type mux struct {
-	conn     packetConn
+	conn     Transport
 	chanList chanList
 
 	incomingChannels chan NewChannel
@@ -113,7 +113,7 @@ func (m *mux) Wait() error {
 }
 
 // newMux returns a mux that runs over the given connection.
-func newMux(p packetConn) *mux {
+func newMux(p Transport) *mux {
 	m := &mux{
 		conn:             p,
 		incomingChannels: make(chan NewChannel, chanSize),
@@ -134,7 +134,7 @@ func (m *mux) sendMessage(msg interface{}) error {
 	if debugMux {
 		log.Printf("send global(%d): %#v", m.chanList.offset, msg)
 	}
-	return m.conn.writePacket(p)
+	return m.conn.WritePacket(p)
 }
 
 func (m *mux) SendRequest(name string, wantReply bool, payload []byte) (bool, []byte, error) {
@@ -212,7 +212,7 @@ func (m *mux) loop() {
 
 // onePacket reads and processes one packet.
 func (m *mux) onePacket() error {
-	packet, err := m.conn.readPacket()
+	packet, err := m.conn.ReadPacket()
 	if err != nil {
 		return err
 	}

ssh/transport.go

@@ -37,6 +37,21 @@ type packetConn interface {
 	Close() error
 }
 
+// Transport represents a connection that implements packet based operations as
+// specified by SSH Transport Protocol (RFC 4253).
+type Transport interface {
+	// WritePacket encrypts and sends a packet of data to the remote peer.
+	WritePacket([]byte) error
+
+	// ReadPacket reads and decrypts a packet of data from the remote peer. The
+	// read is blocking. If error is nil then the returned byte slice is always
+	// non-empty.
+	ReadPacket() ([]byte, error)
+
+	// Close closes the connection with the remote peer.
+	Close() error
+}
+
 // transport is the keyingTransport that implements the SSH packet
 // protocol.
 type transport struct {