package jsonrpc2
import (
"bufio"
"context"
"encoding/json"
"fmt"
"io"
"strconv"
"strings"
"sync"
)
type Stream interface {
Read(context.Context) ([]byte, int64, error)
Write(context.Context, []byte) (int64, error)
}
func NewStream(in io.Reader, out io.Writer) Stream {
return &plainStream{
in: json.NewDecoder(in),
out: out,
}
}
type plainStream struct {
in *json.Decoder
outMu sync.Mutex
out io.Writer
}
func (s *plainStream) Read(ctx context.Context) ([]byte, int64, error) {
select {
case <-ctx.Done():
return nil, 0, ctx.Err()
default:
}
var raw json.RawMessage
if err := s.in.Decode(&raw); err != nil {
return nil, 0, err
}
return raw, int64(len(raw)), nil
}
func (s *plainStream) Write(ctx context.Context, data []byte) (int64, error) {
select {
case <-ctx.Done():
return 0, ctx.Err()
default:
}
s.outMu.Lock()
n, err := s.out.Write(data)
s.outMu.Unlock()
return int64(n), err
}
func NewHeaderStream(in io.Reader, out io.Writer) Stream {
return &headerStream{
in: bufio.NewReader(in),
out: out,
}
}
type headerStream struct {
in *bufio.Reader
outMu sync.Mutex
out io.Writer
}
func (s *headerStream) Read(ctx context.Context) ([]byte, int64, error) {
select {
case <-ctx.Done():
return nil, 0, ctx.Err()
default:
}
var total, length int64
for {
line, err := s.in.ReadString('\n')
total += int64(len(line))
if err != nil {
return nil, total, fmt.Errorf("failed reading header line %q", err)
}
line = strings.TrimSpace(line)
if line == "" {
break
}
colon := strings.IndexRune(line, ':')
if colon < 0 {
return nil, total, fmt.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, total, fmt.Errorf("failed parsing Content-Length: %v", value)
}
if length <= 0 {
return nil, total, fmt.Errorf("invalid Content-Length: %v", length)
}
default:
}
}
if length == 0 {
return nil, total, fmt.Errorf("missing Content-Length header")
}
data := make([]byte, length)
if _, err := io.ReadFull(s.in, data); err != nil {
return nil, total, err
}
total += length
return data, total, nil
}
func (s *headerStream) Write(ctx context.Context, data []byte) (int64, error) {
select {
case <-ctx.Done():
return 0, ctx.Err()
default:
}
s.outMu.Lock()
defer s.outMu.Unlock()
n, err := fmt.Fprintf(s.out, "Content-Length: %v\r\n\r\n", len(data))
total := int64(n)
if err == nil {
n, err = s.out.Write(data)
total += int64(n)
}
return total, err
}