Implement vecnet.ReadFrom on Windows

vecnet/vecnet.go

@@ -31,7 +31,7 @@ type Buffers net.Buffers
 //
 // ReadFrom keeps reading until all bufs are filled or EOF is received.
 //
-// The pre-allocatted space used by ReadFrom is based upon slice lengths.
+// The pre-allocated space used by ReadFrom is based upon slice lengths.
 func (bufs Buffers) ReadFrom(r io.Reader) (int64, error) {
 	if conn, ok := r.(syscall.Conn); ok && readFromBuffers != nil {
 		return readFromBuffers(bufs, conn)

vecnet/vecnet_other.go

@@ -12,7 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// +build !linux linux,386
+//go:build (!linux && !windows) || (linux && 386)
+// +build !linux,!windows linux,386
 
 package vecnet
 

vecnet/vecnet_test.go

@@ -15,7 +15,9 @@
 package vecnet
 
 import (
+	"bytes"
 	"io"
+	"net"
 	"strings"
 	"testing"
 )
@@ -46,3 +48,77 @@ func TestReadFromSanity(t *testing.T) {
 		t.Errorf("ReadFrom() = (%v, %#v), want (%v, %#v)", s1, s2, s[:10], s[10:])
 	}
 }
+
+func TestReadFromNetwork(t *testing.T) {
+	testCloser := func(closer io.Closer) {
+		t.Helper()
+		if err := closer.Close(); err != nil {
+			t.Error(err)
+		}
+	}
+
+	readerConn, readerAddr := newLocalListener(t)
+	defer testCloser(readerConn)
+	writerConn := dialListener(t, readerAddr)
+	defer testCloser(writerConn)
+
+	var (
+		expected = Buffers{
+			[]byte("ABCDEFGHIJKLMNOPQRSTUVWXYZ"),
+			[]byte("0123456789"),
+			[]byte("abcdefghijklmnopqrstuvwxyz"),
+		}
+		received = func() (received Buffers) {
+			received = make(Buffers, len(expected))
+			for i := range received {
+				received[i] = make([]byte, len(expected[i]))
+			}
+			return
+		}()
+	)
+
+	writeMessages(t, writerConn, expected)
+	if _, err := received.ReadFrom(readerConn); err != nil {
+		t.Fatal(err)
+	}
+
+	for i, got := range received {
+		want := expected[i]
+		if !bytes.Equal(got, want) {
+			t.Errorf("buffer %d did not match expected data"+
+				"\n\tgot: %s"+
+				"\n\twant: %s",
+				i, got, want,
+			)
+		}
+	}
+}
+
+func newLocalListener(t *testing.T) (*net.UDPConn, *net.UDPAddr) {
+	t.Helper()
+	addr := &net.UDPAddr{IP: net.IPv4(127, 0, 0, 1)}
+	conn, err := net.ListenUDP("udp", addr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	addr.Port = conn.LocalAddr().(*net.UDPAddr).Port
+	return conn, addr
+}
+
+func dialListener(t *testing.T, readerAddr *net.UDPAddr) *net.UDPConn {
+	t.Helper()
+	writerConn, err := net.DialUDP("udp", nil, readerAddr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	return writerConn
+}
+
+func writeMessages(t *testing.T, conn *net.UDPConn, expected Buffers) {
+	for _, buf := range expected {
+		if _, _, err := conn.WriteMsgUDP(buf, nil, nil); err != nil {
+			t.Fatal(err)
+			return
+		}
+	}
+}

vecnet/vecnet_windows.go

@@ -0,0 +1,95 @@
+// Copyright 2018 The gVisor Authors.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package vecnet
+
+import (
+	"io"
+	"syscall"
+
+	"golang.org/x/sys/windows"
+)
+
+var readFromBuffers = readFromBuffersWindows
+
+func readFromBuffersWindows(bufs Buffers, conn syscall.Conn) (int64, error) {
+	rc, err := conn.SyscallConn()
+	if err != nil {
+		return 0, err
+	}
+
+	length := int64(0)
+	for _, buf := range bufs {
+		length += int64(len(buf))
+	}
+	for n := int64(0); n < length; {
+		cur, err := recvmsg(bufs, rc)
+		if err != nil && (cur == 0 || err != io.EOF) {
+			return n, err
+		}
+		n += int64(cur)
+
+		// Consume buffers to retry.
+		for consumed := 0; consumed < cur; {
+			if len(bufs[0]) <= cur-consumed {
+				consumed += len(bufs[0])
+				bufs = bufs[1:]
+			} else {
+				bufs[0] = bufs[0][cur-consumed:]
+				break
+			}
+		}
+	}
+	return length, nil
+}
+
+func buildWSABufs(bufs Buffers, WSABufs []windows.WSABuf) []windows.WSABuf {
+	for _, buf := range bufs {
+		if l := len(buf); l > 0 {
+			WSABufs = append(WSABufs, windows.WSABuf{
+				Len: uint32(l),
+				Buf: &buf[0],
+			})
+		}
+	}
+	return WSABufs
+}
+
+func recvmsg(bufs Buffers, rc syscall.RawConn) (int, error) {
+	var (
+		bytesReceived uint32
+		WSABufs       = buildWSABufs(bufs, make([]windows.WSABuf, 0, 2))
+		bufCount      = len(bufs)
+
+		msg = windows.WSAMsg{
+			Buffers:     &WSABufs[0],
+			BufferCount: uint32(bufCount),
+		}
+		recvmsgCallBack = func(fd uintptr) bool {
+			winErr := windows.WSARecvMsg(
+				windows.Handle(fd),
+				&msg, &bytesReceived,
+				nil, nil) // TODO: overlapped structure?
+			if winErr != nil {
+				// TODO: double check documentation for other temporary issues
+				// retry if err is temporary
+				canRetry := (winErr == windows.WSAEINTR || winErr == windows.WSAEWOULDBLOCK)
+				return !canRetry
+			}
+			return true
+		}
+		err = rc.Read(recvmsgCallBack)
+	)
+	return int(bytesReceived), err
+}