125 lines
2.9 KiB
Go
125 lines
2.9 KiB
Go
package winpwn
|
|
|
|
import (
|
|
"io"
|
|
"testing"
|
|
"time"
|
|
)
|
|
|
|
// newPipeTube wires a Tube to an in-memory io.Pipe pair instead of a real
|
|
// process/socket, so Tube's framing logic (Send/Recv/RecvUntil/timeouts) can
|
|
// be unit-tested without spawning anything OS-specific. Returns the tube
|
|
// plus the "remote" end the test drives directly.
|
|
func newPipeTube() (tube *Tube, remoteRead *io.PipeReader, remoteWrite *io.PipeWriter) {
|
|
toTube, fromRemote := io.Pipe() // remote writes fromRemote -> tube reads toTube
|
|
toRemote, fromTube := io.Pipe() // tube writes fromTube -> remote reads toRemote
|
|
tube = newTube(nil, nil, fromTube, toTube)
|
|
return tube, toRemote, fromRemote
|
|
}
|
|
|
|
func TestTubeSendRecv(t *testing.T) {
|
|
tube, remoteRead, remoteWrite := newPipeTube()
|
|
defer tube.Close()
|
|
|
|
go func() {
|
|
buf := make([]byte, 5)
|
|
_, _ = io.ReadFull(remoteRead, buf)
|
|
_, _ = remoteWrite.Write([]byte("got: "))
|
|
_, _ = remoteWrite.Write(buf)
|
|
}()
|
|
|
|
if err := tube.Send([]byte("hello")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
|
|
got, err := tube.Recv(10)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if string(got) != "got: hello" {
|
|
t.Errorf("Recv = %q, want %q", got, "got: hello")
|
|
}
|
|
}
|
|
|
|
func TestTubeRecvUntil(t *testing.T) {
|
|
tube, _, remoteWrite := newPipeTube()
|
|
defer tube.Close()
|
|
|
|
go func() {
|
|
_, _ = remoteWrite.Write([]byte("foo bar: baz\n"))
|
|
}()
|
|
|
|
got, err := tube.RecvUntil([]byte(": "))
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if string(got) != "foo bar: " {
|
|
t.Errorf("RecvUntil = %q, want %q", got, "foo bar: ")
|
|
}
|
|
}
|
|
|
|
func TestTubeRecvLineUsesContextNewline(t *testing.T) {
|
|
tube, _, remoteWrite := newPipeTube()
|
|
defer tube.Close()
|
|
|
|
go func() {
|
|
_, _ = remoteWrite.Write([]byte("first line\nsecond"))
|
|
}()
|
|
|
|
got, err := tube.RecvLine()
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if string(got) != "first line\n" {
|
|
t.Errorf("RecvLine = %q, want %q", got, "first line\n")
|
|
}
|
|
}
|
|
|
|
func TestTubeSendAfter(t *testing.T) {
|
|
tube, remoteRead, remoteWrite := newPipeTube()
|
|
defer tube.Close()
|
|
|
|
go func() {
|
|
_, _ = remoteWrite.Write([]byte("password: "))
|
|
buf := make([]byte, 6)
|
|
_, _ = io.ReadFull(remoteRead, buf)
|
|
_, _ = remoteWrite.Write(buf)
|
|
}()
|
|
|
|
if err := tube.SendAfter([]byte(": "), []byte("secret")); err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
got, err := tube.Recv(6)
|
|
if err != nil {
|
|
t.Fatal(err)
|
|
}
|
|
if string(got) != "secret" {
|
|
t.Errorf("echoed payload = %q, want %q", got, "secret")
|
|
}
|
|
}
|
|
|
|
func TestTubeRecvTimesOutOnHang(t *testing.T) {
|
|
tube, _, _ := newPipeTube() // nobody ever writes
|
|
defer tube.Close()
|
|
tube.SetTimeout(50 * time.Millisecond)
|
|
|
|
start := time.Now()
|
|
_, err := tube.Recv(1)
|
|
if err == nil {
|
|
t.Fatal("expected a timeout error, got nil")
|
|
}
|
|
if elapsed := time.Since(start); elapsed > time.Second {
|
|
t.Errorf("timeout took %s, expected ~50ms", elapsed)
|
|
}
|
|
}
|
|
|
|
func TestTubeCloseIsIdempotent(t *testing.T) {
|
|
tube, _, _ := newPipeTube()
|
|
if err := tube.Close(); err != nil {
|
|
t.Fatalf("first Close: %v", err)
|
|
}
|
|
if err := tube.Close(); err != nil {
|
|
t.Fatalf("second Close should be a no-op, got: %v", err)
|
|
}
|
|
}
|