msg: Return usable object after handshake

This commit is contained in:
Matt Joiner 2015-03-12 20:29:48 +11:00
parent b3c4afbe25
commit 13a5b8b279
2 changed files with 18 additions and 4 deletions

View File

@ -339,7 +339,12 @@ func readUntil(r io.Reader, b []byte) error {
return nil
}
func (h *handshake) Do() (ret io.ReadWriteCloser, err error) {
type readWriter struct {
io.Reader
io.Writer
}
func (h *handshake) Do() (ret io.ReadWriter, err error) {
err = h.establishS()
if err != nil {
return
@ -383,6 +388,7 @@ func (h *handshake) Do() (ret io.ReadWriteCloser, err error) {
err = fmt.Errorf("error reading crypto negotiation: %s", err)
return
}
ret = readWriter{r, &cipherWriter{bC, h.conn}}
} else {
err = readUntil(h.conn, hash(req1, h.s.Bytes()))
if err != nil {
@ -422,6 +428,7 @@ func (h *handshake) Do() (ret io.ReadWriteCloser, err error) {
if err != nil {
return
}
ret = readWriter{r, w}
}
err = h.finishWriting()
if err != nil {
@ -431,7 +438,7 @@ func (h *handshake) Do() (ret io.ReadWriteCloser, err error) {
return
}
func Handshake(rw io.ReadWriteCloser, initer bool, skey []byte) (ret io.ReadWriteCloser, err error) {
func Handshake(rw io.ReadWriteCloser, initer bool, skey []byte) (ret io.ReadWriter, err error) {
h := handshake{
conn: rw,
initer: initer,

View File

@ -3,6 +3,7 @@ package mse
import (
"bytes"
"io"
"log"
"net"
"sync"
@ -53,7 +54,10 @@ func TestHandshake(t *testing.T) {
t.Fatal(err)
return
}
a.Close()
a.Write([]byte("hello world"))
var msg [20]byte
n, _ := a.Read(msg[:])
log.Print(string(msg[:n]))
}()
go func() {
defer wg.Done()
@ -62,7 +66,10 @@ func TestHandshake(t *testing.T) {
t.Fatal(err)
return
}
b.Close()
var msg [20]byte
n, _ := b.Read(msg[:])
log.Print(string(msg[:n]))
b.Write([]byte("yo dawg"))
}()
wg.Wait()
}