From 3eabd1969d018c0c61703fad4db0708f86ff97e3 Mon Sep 17 00:00:00 2001
From: John Starks <jostarks@microsoft.com>
Date: Mon, 25 Jun 2018 09:31:55 -0700
Subject: [PATCH] Clean up message mode pipe handling, add test

---
 file.go      |  3 ---
 pipe.go      | 18 +++++----------
 pipe_test.go | 63 ++++++++++++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 68 insertions(+), 16 deletions(-)

diff --git a/file.go b/file.go
index 3c0866a..4334ff1 100644
--- a/file.go
+++ b/file.go
@@ -229,9 +229,6 @@ func (f *win32File) Read(b []byte) (int, error) {
 		return 0, io.EOF
 	} else if err == syscall.ERROR_BROKEN_PIPE {
 		return 0, io.EOF
-	// When there is more data in the message pipe to read, we get ERROR_MORE_DATA. We ignore that error and proceed to read more data
-	} else if err == syscall.ERROR_MORE_DATA && n != 0 && len(b) != 0 {
-		return n, nil
 	} else {
 		return n, err
 	}
diff --git a/pipe.go b/pipe.go
index 34c63e7..806fd34 100644
--- a/pipe.go
+++ b/pipe.go
@@ -121,6 +121,11 @@ func (f *win32MessageBytePipe) Read(b []byte) (int, error) {
 		// zero-byte message, ensure that all future Read() calls
 		// also return EOF.
 		f.readEOF = true
+	} else if err == syscall.ERROR_MORE_DATA {
+		// ERROR_MORE_DATA indicates that the pipe's read mode is message mode
+		// and the message still has more bytes. Treat this as a success, since
+		// this package presents all named pipes as byte streams.
+		err = nil
 	}
 	return n, err
 }
@@ -175,19 +180,6 @@ func DialPipe(path string, timeout *time.Duration) (net.Conn, error) {
 		return nil, err
 	}
 
-	var state uint32
-	err = getNamedPipeHandleState(h, &state, nil, nil, nil, nil, 0)
-	if err != nil {
-		return nil, err
-	}
-
-	/**
-	Windows support message type pipes in message-read mode only. Removing this check to allow for windows named pipes.
-	*/
-	/*if state&cPIPE_READMODE_MESSAGE != 0 {
-		return nil, &os.PathError{Op: "open", Path: path, Err: errors.New("message readmode pipes not supported")}
-	}*/
-
 	f, err := makeWin32File(h)
 	if err != nil {
 		syscall.Close(h)
diff --git a/pipe_test.go b/pipe_test.go
index c0d1a77..3869207 100644
--- a/pipe_test.go
+++ b/pipe_test.go
@@ -2,12 +2,15 @@ package winio
 
 import (
 	"bufio"
+	"bytes"
 	"io"
 	"net"
 	"os"
+	"sync"
 	"syscall"
 	"testing"
 	"time"
+	"unsafe"
 )
 
 var testPipeName = `\\.\pipe\winiotestpipe`
@@ -451,3 +454,63 @@ func TestConnectRace(t *testing.T) {
 		c.Close()
 	}
 }
+
+func TestMessageReadMode(t *testing.T) {
+	var wg sync.WaitGroup
+	defer wg.Wait()
+
+	l, err := ListenPipe(testPipeName, &PipeConfig{MessageMode: true})
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer l.Close()
+
+	msg := ([]byte)("hello world")
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		s, err := l.Accept()
+		if err != nil {
+			t.Fatal(err)
+		}
+		_, err = s.Write(msg)
+		if err != nil {
+			t.Fatal(err)
+		}
+		s.Close()
+	}()
+
+	c, err := DialPipe(testPipeName, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer c.Close()
+
+	setNamedPipeHandleState := syscall.NewLazyDLL("kernel32.dll").NewProc("SetNamedPipeHandleState")
+
+	p := c.(*win32MessageBytePipe)
+	mode := uint32(cPIPE_READMODE_MESSAGE)
+	if s, _, err := setNamedPipeHandleState.Call(uintptr(p.handle), uintptr(unsafe.Pointer(&mode)), 0, 0); s == 0 {
+		t.Fatal(err)
+	}
+
+	ch := make([]byte, 1)
+	var vmsg []byte
+	for {
+		n, err := c.Read(ch)
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			t.Fatal(err)
+		}
+		if n != 1 {
+			t.Fatal("expected 1: ", n)
+		}
+		vmsg = append(vmsg, ch[0])
+	}
+	if !bytes.Equal(msg, vmsg) {
+		t.Fatalf("expected %s: %s", msg, vmsg)
+	}
+}