/
stream.go
161 lines (149 loc) · 3.98 KB
/
stream.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.
package jsonrpc2
import (
"bufio"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"strconv"
"strings"
"sync"
)
// Stream abstracts the transport mechanics from the JSON RPC protocol.
// A Conn reads and writes messages using the stream it was provided on
// construction, and assumes that each call to Read or Write fully transfers
// a single message, or returns an error.
type Stream interface {
// Read gets the next message from the stream.
// It is never called concurrently.
Read(context.Context) ([]byte, int64, error)
// Write sends a message to the stream.
// It must be safe for concurrent use.
Write(context.Context, []byte) (int64, error)
}
// ErrDisconnected signals that the stream or connection exited normally.
var ErrDisconnected = errors.New("disconnected")
// NewStream returns a Stream built on top of an io.Reader and io.Writer
// The messages are sent with no wrapping, and rely on json decode consistency
// to determine message boundaries.
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 {
if err == io.EOF {
return nil, 0, ErrDisconnected
}
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
}
// NewHeaderStream returns a Stream built on top of an io.Reader and io.Writer
// The messages are sent with HTTP content length and MIME type headers.
// This is the format used by LSP and others.
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
// read the header, stop on the first empty line
for {
line, err := s.in.ReadString('\n')
total += int64(len(line))
if err == io.EOF {
// A normal disconnection will terminate with EOF before the next header.
return nil, total, ErrDisconnected
}
if err != nil {
return nil, total, fmt.Errorf("failed reading header line %q", err)
}
line = strings.TrimSpace(line)
// check we have a header 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:
// ignoring unknown headers
}
}
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
}