Add ByteStreamWrapper to transparently copying ReadWriteClosers

Signed-off-by: Derek McGowan <derek@mcgstyle.net> (github: dmcgowan)

Author: dmcgowan

inmem.go

@@ -156,6 +156,33 @@ func (s *streamSession) decodeStream(v reflect.Value, b []byte) error {
 	return nil
 }
 
+func (s *streamSession) encodeWrapper(v reflect.Value) ([]byte, error) {
+	wrapper := v.Interface().(ByteStreamWrapper)
+	bs, err := s.newByteStream()
+	if err != nil {
+		return nil, err
+	}
+
+	go func() {
+		io.Copy(bs, wrapper)
+		bs.Close()
+	}()
+
+	go func() {
+		io.Copy(wrapper, bs)
+		wrapper.Close()
+	}()
+
+	return s.encodeStream(reflect.ValueOf(bs).Elem())
+}
+
+func (s *streamSession) decodeWrapper(v reflect.Value, b []byte) error {
+	bs := &byteStream{}
+	s.decodeStream(reflect.ValueOf(bs).Elem(), b)
+	v.FieldByName("ReadWriteCloser").Set(reflect.ValueOf(bs))
+	return nil
+}
+
 func getMsgPackHandler(session *streamSession) *codec.MsgpackHandle {
 	mh := &codec.MsgpackHandle{WriteExt: true}
 	mh.RawToString = true
@@ -175,6 +202,11 @@ func getMsgPackHandler(session *streamSession) *codec.MsgpackHandle {
 		panic(err)
 	}
 
+	err = mh.AddExt(reflect.TypeOf(ByteStreamWrapper{}), 4, session.encodeWrapper, session.decodeWrapper)
+	if err != nil {
+		panic(err)
+	}
+
 	return mh
 }
 

inmem_test.go

@@ -303,6 +303,44 @@ func TestComplexMessage(t *testing.T) {
 	SpawnPipeTestRoutines(t, client, server)
 }
 
+func TestInmemWrappedSend(t *testing.T) {
+	tmp, err := ioutil.TempFile("", "libchan-test-")
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer os.RemoveAll(tmp.Name())
+	fmt.Fprintf(tmp, "hello through a wrapper\n")
+	tmp.Sync()
+	tmp.Seek(0, 0)
+
+	client := func(t *testing.T, w Sender) {
+		message := &InMemMessage{Data: "path=" + tmp.Name(), Stream: ByteStreamWrapper{tmp}}
+		err = w.Send(message)
+		if err != nil {
+			t.Fatal(err)
+		}
+	}
+	server := func(t *testing.T, r Receiver) {
+		msg := &InMemMessage{}
+		err := r.Receive(msg)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if msg.Data != "path="+tmp.Name() {
+			t.Fatalf("%#v", msg)
+		}
+		txt, err := ioutil.ReadAll(msg.Stream)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if string(txt) != "hello through a wrapper\n" {
+			t.Fatalf("%s\n", txt)
+		}
+
+	}
+	SpawnPipeTestRoutines(t, client, server)
+}
+
 type SendTestRoutine func(*testing.T, Sender)
 type ReceiveTestRoutine func(*testing.T, Receiver)
 

libchan.go

@@ -60,3 +60,12 @@ type Receiver interface {
 	// automatically closed through receiving an EOF.
 	Close() error
 }
+
+// ByteStreamWrapper is a wrapper around a ReadWriteCloser
+// to cue the transport to copy to a transport byte stream.
+// Note: ReadWriteClosers created through calling the
+// CreateByteStream method on a Sender do not need
+// to wrap the ByteStream.
+type ByteStreamWrapper struct {
+	io.ReadWriteCloser
+}

spdy/encode.go

@@ -3,11 +3,13 @@ package spdy
 import (
 	"encoding/binary"
 	"errors"
+	"io"
 	"net"
 	"reflect"
 	"time"
 
 	"github.com/dmcgowan/go/codec"
+	"github.com/docker/libchan"
 )
 
 func (s *SpdyTransport) encodeChannel(v reflect.Value) ([]byte, error) {
@@ -87,6 +89,33 @@ func (s *SpdyTransport) decodeStream(v reflect.Value, b []byte) error {
 	return nil
 }
 
+func (s *SpdyTransport) encodeWrapper(v reflect.Value) ([]byte, error) {
+	wrapper := v.Interface().(libchan.ByteStreamWrapper)
+	bs, err := s.createByteStream()
+	if err != nil {
+		return nil, err
+	}
+
+	go func() {
+		io.Copy(bs, wrapper)
+		bs.Close()
+	}()
+
+	go func() {
+		io.Copy(wrapper, bs)
+		wrapper.Close()
+	}()
+
+	return s.encodeStream(reflect.ValueOf(bs).Elem())
+}
+
+func (s *SpdyTransport) decodeWrapper(v reflect.Value, b []byte) error {
+	bs := &byteStream{}
+	s.decodeStream(reflect.ValueOf(bs).Elem(), b)
+	v.FieldByName("ReadWriteCloser").Set(reflect.ValueOf(bs))
+	return nil
+}
+
 func (s *SpdyTransport) waitConn(network, local, remote string, timeout time.Duration) (net.Conn, error) {
 	timeoutChan := time.After(timeout)
 	connChan := make(chan net.Conn)
@@ -218,6 +247,11 @@ func (s *SpdyTransport) initializeHandler() *codec.MsgpackHandle {
 		panic(err)
 	}
 
+	err = mh.AddExt(reflect.TypeOf(libchan.ByteStreamWrapper{}), 0x03, s.encodeWrapper, s.decodeWrapper)
+	if err != nil {
+		panic(err)
+	}
+
 	// Register networks
 	s.networks["tcp"] = 0x04
 	s.netConns[0x04] = make(map[string]net.Conn)

spdy/session_test.go

@@ -292,6 +292,74 @@ func TestByteStream(t *testing.T) {
 	SpawnClientServerTest(t, "localhost:12944", ClientSendWrapper(client), ServerReceiveWrapper(server))
 }
 
+type WrappedMessage struct {
+	Message string
+	Wrapped io.ReadWriteCloser
+}
+
+func TestWrappedByteStreams(t *testing.T) {
+	serverSend := "G'day client ☺"
+	clientReply := "Hello Server, ☢ FYI your stream was transparently copied ☠"
+	client := func(t *testing.T, sender libchan.Sender, s *SpdyTransport) {
+		// Create pipe
+		p1, p2 := net.Pipe()
+
+		m1 := &WrappedMessage{
+			Message: "wrapped",
+			Wrapped: libchan.ByteStreamWrapper{p2},
+		}
+
+		sendErr := sender.Send(m1)
+		if sendErr != nil {
+			t.Fatalf("Error sending channel: %s", sendErr)
+		}
+
+		// read
+		readBytes := make([]byte, 30)
+		n, readErr := p1.Read(readBytes)
+		if readErr != nil {
+			t.Fatalf("Error reading from byte stream: %s", readErr)
+		}
+		if expected := serverSend; string(readBytes[:n]) != expected {
+			t.Fatalf("Unexpected read value:\n\tExpected: %q\n\tActual: %q", expected, string(readBytes[:n]))
+		}
+
+		// write
+		_, writeErr := p1.Write([]byte(clientReply))
+		if writeErr != nil {
+			t.Fatalf("Error writing to byte stream: %s", writeErr)
+		}
+
+	}
+	server := func(t *testing.T, receiver libchan.Receiver, s *SpdyTransport) {
+		m1 := &WrappedMessage{}
+		recvErr := receiver.Receive(m1)
+		if recvErr != nil {
+			t.Fatalf("Error receiving message: %s", recvErr)
+		}
+
+		if expected := "wrapped"; m1.Message != expected {
+			t.Fatalf("Unexpected message\n\tExpected: %s\n\tActual: %s", expected, m1.Message)
+		}
+
+		_, writeErr := m1.Wrapped.Write([]byte(serverSend))
+		if writeErr != nil {
+			t.Fatalf("Error writing to byte stream: %s", writeErr)
+		}
+
+		readBytes := make([]byte, 80)
+		n, readErr := m1.Wrapped.Read(readBytes)
+		if readErr != nil {
+			t.Fatalf("Error reading from byte stream: %s", readErr)
+		}
+		if expected := clientReply; string(readBytes[:n]) != expected {
+			t.Fatalf("Unexpected read value:\n\tExpected: %q\n\tActual: %q", expected, string(readBytes[:n]))
+		}
+
+	}
+	SpawnClientServerTest(t, "localhost:12943", ClientSendWrapper(client), ServerReceiveWrapper(server))
+}
+
 func ClientSendWrapper(f func(t *testing.T, c libchan.Sender, s *SpdyTransport)) ClientRoutine {
 	return func(t *testing.T, server string) {
 		conn, connErr := net.Dial("tcp", server)