Update interface to make all byte streams transparent

All ReadWriteClosers and channels will now be copied on send unless detected to have been created locally.


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

Author: dmcgowan

copy.go

@@ -19,77 +19,12 @@ func Copy(w Sender, r Receiver) (int, error) {
 				return n, err
 			}
 		}
-		mCopy, err := copyChannelMessage(w, m)
-		if err != nil {
-			return n, err
-		}
 
-		err = w.Send(mCopy)
+		err = w.Send(m)
 		if err != nil {
 			return n, err
 		}
 		n++
 	}
 	return n, nil
 }
-
-func copyByteStream(sender Sender, stream io.ReadWriteCloser) (io.ReadWriteCloser, error) {
-	streamCopy, err := sender.CreateByteStream()
-	if err != nil {
-		return nil, err
-	}
-	go func() {
-		io.Copy(streamCopy, stream)
-		streamCopy.Close()
-	}()
-	go func() {
-		io.Copy(stream, streamCopy)
-		stream.Close()
-	}()
-	return streamCopy, nil
-}
-
-func copyChannelMessage(sender Sender, m map[string]interface{}) (map[string]interface{}, error) {
-	mCopy := make(map[string]interface{})
-	for k, v := range m {
-		// Throw error if tcp/udp connections?
-		switch val := v.(type) {
-		case io.ReadWriteCloser:
-			streamCopy, err := copyByteStream(sender, val)
-			if err != nil {
-				return nil, err
-			}
-			mCopy[k] = streamCopy
-		case Sender:
-			recv, send, err := sender.CreateNestedReceiver()
-			if err != nil {
-				return nil, err
-			}
-			go func() {
-				Copy(val, recv)
-				// TODO propagate or log error
-			}()
-			mCopy[k] = send
-		case Receiver:
-			send, recv, err := sender.CreateNestedSender()
-			if err != nil {
-				return nil, err
-			}
-			go func() {
-				Copy(send, val)
-				// TODO propagate or log error
-			}()
-			mCopy[k] = recv
-		case map[string]interface{}:
-			vCopy, err := copyChannelMessage(sender, val)
-			if err != nil {
-				return nil, err
-			}
-			mCopy[k] = vCopy
-		default:
-			mCopy[k] = v
-		}
-	}
-
-	return mCopy, nil
-}

copy_test.go

@@ -2,6 +2,7 @@ package libchan
 
 import (
 	"io"
+	"net"
 	"os"
 	"runtime/pprof"
 	"testing"
@@ -88,16 +89,13 @@ func TestByteStreamProxy(t *testing.T) {
 	sendString := "Sending a string"
 	retString := "Returned string"
 	client := func(t *testing.T, sender Sender) {
-		bs, err := sender.CreateByteStream()
-		if err != nil {
-			t.Fatalf("Error creating byte stream: %s", err)
-		}
+		bs, bsRemote := net.Pipe()
 
 		message := &ProxiedStreamMessage{
-			Stream: bs,
+			Stream: bsRemote,
 		}
 
-		err = sender.Send(message)
+		err := sender.Send(message)
 		if err != nil {
 			t.Fatalf("Error sending message: %s", err)
 		}

inmem.go

@@ -60,6 +60,7 @@ func (s *streamSession) newByteStream() (io.ReadWriteCloser, error) {
 	bs := &byteStream{
 		Conn:        c1,
 		referenceID: s.referenceID,
+		session:     s,
 	}
 	s.referenceLock.Lock()
 	s.byteStreams[s.referenceID] = bs
@@ -157,53 +158,62 @@ 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
-	}
+// Wrappers to cue a copy of a previously decoded object
 
-	go func() {
-		io.Copy(bs, wrapper)
-		bs.Close()
-	}()
+type pipeSenderWrapper struct{ *pipeSender }
+type pipeReceiverWrapper struct{ *pipeReceiver }
+type byteStreamWrapper struct{ *byteStream }
 
-	go func() {
-		io.Copy(wrapper, bs)
-		wrapper.Close()
-	}()
+func (s *streamSession) decodeSenderWrapper(v reflect.Value, b []byte) error {
+	ps := &pipeSender{}
+	v.FieldByName("pipeSender").Set(reflect.ValueOf(ps))
+	return s.decodeSender(reflect.ValueOf(ps).Elem(), b)
+}
 
-	return s.encodeStream(reflect.ValueOf(bs).Elem())
+func (s *streamSession) decodeReceiverWrapper(v reflect.Value, b []byte) error {
+	pr := &pipeReceiver{}
+	v.FieldByName("pipeReceiver").Set(reflect.ValueOf(pr))
+	return s.decodeReceiver(reflect.ValueOf(pr).Elem(), b)
 }
 
-func (s *streamSession) decodeWrapper(v reflect.Value, b []byte) error {
+func (s *streamSession) decodeByteStreamWrapper(v reflect.Value, b []byte) error {
 	bs := &byteStream{}
-	s.decodeStream(reflect.ValueOf(bs).Elem(), b)
-	v.FieldByName("ReadWriteCloser").Set(reflect.ValueOf(bs))
-	return nil
+	v.FieldByName("byteStream").Set(reflect.ValueOf(bs))
+	return s.decodeStream(reflect.ValueOf(bs).Elem(), b)
 }
 
+// Internal definitions early, does not follow protocol because does not
+// interact with external processes
 func getMsgPackHandler(session *streamSession) *codec.MsgpackHandle {
 	mh := &codec.MsgpackHandle{WriteExt: true}
 	mh.RawToString = true
 
-	err := mh.AddExt(reflect.TypeOf(pipeReceiver{}), 1, session.encodeReceiver, session.decodeReceiver)
+	err := mh.AddExt(reflect.TypeOf(pipeReceiverWrapper{}), 1, nil, session.decodeReceiverWrapper)
 	if err != nil {
 		panic(err)
 	}
 
-	err = mh.AddExt(reflect.TypeOf(pipeSender{}), 2, session.encodeSender, session.decodeSender)
+	err = mh.AddExt(reflect.TypeOf(pipeReceiver{}), 1, session.encodeReceiver, nil)
 	if err != nil {
 		panic(err)
 	}
 
-	err = mh.AddExt(reflect.TypeOf(byteStream{}), 3, session.encodeStream, session.decodeStream)
+	err = mh.AddExt(reflect.TypeOf(pipeSenderWrapper{}), 2, nil, session.decodeSenderWrapper)
 	if err != nil {
 		panic(err)
 	}
 
-	err = mh.AddExt(reflect.TypeOf(ByteStreamWrapper{}), 4, session.encodeWrapper, session.decodeWrapper)
+	err = mh.AddExt(reflect.TypeOf(pipeSender{}), 2, session.encodeSender, nil)
+	if err != nil {
+		panic(err)
+	}
+
+	err = mh.AddExt(reflect.TypeOf(byteStreamWrapper{}), 3, nil, session.decodeByteStreamWrapper)
+	if err != nil {
+		panic(err)
+	}
+
+	err = mh.AddExt(reflect.TypeOf(byteStream{}), 3, session.encodeStream, nil)
 	if err != nil {
 		panic(err)
 	}
@@ -219,7 +229,11 @@ type pipeSender struct {
 }
 
 func (w *pipeSender) Send(message interface{}) error {
-	return w.encoder.Encode(message)
+	mCopy, mErr := w.copyMessage(message)
+	if mErr != nil {
+		return mErr
+	}
+	return w.encoder.Encode(mCopy)
 }
 
 func (w *pipeSender) Close() error {
@@ -261,3 +275,115 @@ type byteStream struct {
 	referenceID uint64
 	session     *streamSession
 }
+
+func (w *pipeSender) copyMessage(message interface{}) (interface{}, error) {
+	mapCopy, mapOk := message.(map[string]interface{})
+	if mapOk {
+		return w.copyChannelMessage(mapCopy)
+	}
+	return w.copyStructure(message)
+}
+
+func (w *pipeSender) copyValue(v interface{}) (interface{}, error) {
+	switch val := v.(type) {
+	case *byteStream:
+		if val.session != w.session {
+			return w.copyByteStream(val)
+		}
+	case *pipeSender:
+		if val.session != w.session {
+			return w.copySender(val)
+		}
+	case *pipeReceiver:
+		if val.session != w.session {
+			return w.copyReceiver(val)
+		}
+	case io.ReadWriteCloser:
+		return w.copyByteStream(val)
+	case Sender:
+		return w.copySender(val)
+	case Receiver:
+		return w.copyReceiver(val)
+	case map[string]interface{}:
+		return w.copyChannelMessage(val)
+	case struct{}:
+		return w.copyStructure(v)
+	case *struct{}:
+		return w.copyStructure(v)
+	default:
+	}
+	return v, nil
+}
+
+func (w *pipeSender) copySender(val Sender) (Sender, error) {
+	recv, send, err := w.CreateNestedReceiver()
+	if err != nil {
+		return nil, err
+	}
+	go func() {
+		Copy(val, recv)
+		val.Close()
+	}()
+	return send, nil
+}
+
+func (w *pipeSender) copyReceiver(val Receiver) (Receiver, error) {
+	send, recv, err := w.CreateNestedSender()
+	if err != nil {
+		return nil, err
+	}
+	go func() {
+		Copy(send, val)
+		send.Close()
+	}()
+	return recv, nil
+}
+
+func (w *pipeSender) copyByteStream(stream io.ReadWriteCloser) (io.ReadWriteCloser, error) {
+	streamCopy, err := w.session.newByteStream()
+	if err != nil {
+		return nil, err
+	}
+	go func() {
+		io.Copy(streamCopy, stream)
+		streamCopy.Close()
+	}()
+	go func() {
+		io.Copy(stream, streamCopy)
+		stream.Close()
+	}()
+	return streamCopy, nil
+}
+
+func (w *pipeSender) copyChannelMessage(m map[string]interface{}) (interface{}, error) {
+	mCopy := make(map[string]interface{})
+	for k, v := range m {
+		vCopy, vErr := w.copyValue(v)
+		if vErr != nil {
+			return nil, vErr
+		}
+		mCopy[k] = vCopy
+	}
+
+	return mCopy, nil
+}
+
+func (w *pipeSender) copyStructure(m interface{}) (interface{}, error) {
+	v := reflect.ValueOf(m)
+	if v.Kind() == reflect.Ptr {
+		v = v.Elem()
+	}
+	if v.Kind() != reflect.Struct {
+		return nil, errors.New("invalid non struct type")
+	}
+	mCopy := make(map[string]interface{})
+	t := v.Type()
+	for i := 0; i < v.NumField(); i++ {
+		vCopy, vErr := w.copyValue(v.Field(i).Interface())
+		if vErr != nil {
+			return nil, vErr
+		}
+		mCopy[t.Field(i).Name] = vCopy
+	}
+	return mCopy, nil
+}