stream: chnage Stream interface methods arg to like io.ReadWriter

Author: zchee

stream.go

@@ -19,9 +19,10 @@ import (
 // Stream abstracts the transport mechanics from the JSON RPC protocol.
 type Stream interface {
 	// Read gets the next message from the stream.
-	Read(context.Context) ([]byte, error)
+	Read(ctx context.Context, p []byte) (n int, err error)
+
 	// Write sends a message to the stream.
-	Write(context.Context, []byte) error
+	Write(ctx context.Context, p []byte) (n int, err error)
 }
 
 type stream struct {
@@ -30,26 +31,25 @@ type stream struct {
 	sync.Mutex
 }
 
-// NewStream returns the new Stream.
 func NewStream(in io.Reader, out io.Writer) Stream {
 	return &stream{
 		in:  bufio.NewReader(in),
 		out: out,
 	}
 }
 
-func (s *stream) Read(ctx context.Context) ([]byte, error) {
+func (s *stream) Read(ctx context.Context, p []byte) (n int, err error) {
 	select {
 	case <-ctx.Done():
-		return nil, ctx.Err()
+		return 0, ctx.Err()
 	default:
 	}
 
 	var length int64
 	for {
 		line, err := s.in.ReadString('\n')
 		if err != nil {
-			return nil, xerrors.Errorf("failed reading header line: %w", err)
+			return 0, xerrors.Errorf("failed reading header line: %w", err)
 		}
 
 		line = strings.TrimSpace(line)
@@ -59,49 +59,49 @@ func (s *stream) Read(ctx context.Context) ([]byte, error) {
 
 		colon := strings.IndexRune(line, ':')
 		if colon < 0 {
-			return nil, xerrors.Errorf("invalid header line: %q", line)
+			return 0, xerrors.Errorf("invalid header line: %q", line)
 		}
 
 		name, value := line[:colon], strings.TrimSpace(line[colon+1:])
-		switch name {
-		case "Content-Length":
-			if length, err = strconv.ParseInt(value, 10, 32); err != nil {
-				return nil, xerrors.Errorf("failed parsing Content-Length: %v", value)
-			}
-
-			if length <= 0 {
-				return nil, xerrors.Errorf("invalid Content-Length: %v", length)
-			}
-		default:
-			// ignoring unknown headers
+		if name != "Content-Length" {
+			continue
+		}
+
+		if length, err = strconv.ParseInt(value, 10, 32); err != nil {
+			return 0, xerrors.Errorf("failed parsing Content-Length: %v", value)
+		}
+
+		if length <= 0 {
+			return 0, xerrors.Errorf("invalid Content-Length: %v", length)
 		}
 	}
 
 	if length == 0 {
-		return nil, xerrors.New("missing Content-Length header")
+		return 0, xerrors.New("missing Content-Length header")
 	}
 
-	data := make([]byte, length)
-	if _, err := io.ReadFull(s.in, data); err != nil {
-		return nil, xerrors.Errorf("failed reading data: %w", err)
+	p = make([]byte, length)
+	n, err = io.ReadFull(s.in, p)
+	if err != nil {
+		return 0, xerrors.Errorf("failed reading data: %w", err)
 	}
 
-	return data, nil
+	return n, nil
 }
 
-func (s *stream) Write(ctx context.Context, data []byte) error {
+func (s *stream) Write(ctx context.Context, p []byte) (n int, err error) {
 	select {
 	case <-ctx.Done():
-		return ctx.Err()
+		return 0, ctx.Err()
 	default:
 	}
 
 	s.Lock()
-	_, err := fmt.Fprintf(s.out, "Content-Length: %v\r\n\r\n", len(data))
+	n, err = fmt.Fprintf(s.out, "Content-Length: %v\r\n\r\n", len(p))
 	if err == nil {
-		_, err = s.out.Write(data)
+		n, err = s.out.Write(p)
 	}
 	s.Unlock()
 
-	return err
+	return n, err
 }