forked from sourcegraph/httpfstream
-
Notifications
You must be signed in to change notification settings - Fork 0
/
client.go
158 lines (141 loc) · 3.45 KB
/
client.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
package httpfstream
import (
"bytes"
"crypto/tls"
"errors"
"fmt"
"github.com/garyburd/go-websocket/websocket"
"io"
"net"
"net/http"
"net/url"
"os"
"strings"
"time"
)
// Follow opens a WebSocket to the file at the given URL (which must be handled
// by httpfstream's HTTP handler) and returns the file's contents. The
// io.ReadCloser continues to return data (blocking as needed) if, and as long
// as, there is an active writer to the file.
func Follow(u *url.URL) (io.ReadCloser, error) {
ws, resp, err := newClient(u, "FOLLOW")
if err == websocket.ErrBadHandshake {
err = errorFromResponse(resp, nil)
}
if err != nil {
return nil, err
}
if resp.StatusCode == http.StatusOK {
return resp.Body, nil
}
return &webSocketReadCloser{ws}, nil
}
type webSocketReadCloser struct {
ws *websocket.Conn
}
// Read implements io.Reader.
func (r *webSocketReadCloser) Read(p []byte) (n int, err error) {
op, rdr, err := r.ws.NextReader()
if err != nil {
return 0, err
}
if op != websocket.OpText {
return 0, errors.New("websocket op is not text")
}
n, err = rdr.Read(p)
if err == io.EOF {
return r.Read(p)
}
return
}
// Close implements io.Closer.
func (r *webSocketReadCloser) Close() error {
return r.ws.Close()
}
// Append appends data from r to the file at the given URL.
func Append(u *url.URL, r io.Reader) error {
w, err := OpenAppend(u)
if err != nil {
return err
}
defer w.Close()
_, err = io.Copy(w, r)
return err
}
// OpenAppend opens a WebSocket to the file at the given URL (which must point
// be handled by httpfstream's HTTP handler) and returns an io.WriteCloser that writes
// (via the WebSocket) to that file.
func OpenAppend(u *url.URL) (io.WriteCloser, error) {
ws, resp, err := newClient(u, "APPEND")
if resp != nil {
defer resp.Body.Close()
}
if err != nil {
if err == websocket.ErrBadHandshake {
err2 := errorFromResponse(resp, nil)
if err2 != nil {
return nil, err2
}
return nil, err
}
return nil, err
}
return &appendWriteCloser{new(bytes.Buffer), ws}, nil
}
type appendWriteCloser struct {
io.Writer
ws *websocket.Conn
}
// Write implements io.Writer.
func (pw *appendWriteCloser) Write(p []byte) (n int, err error) {
pw.ws.SetWriteDeadline(time.Now().Add(writeWait))
w, err := pw.ws.NextWriter(websocket.OpText)
if err != nil {
return 0, err
}
defer w.Close()
return w.Write(p)
}
// Write implements io.Closer.
func (pw *appendWriteCloser) Close() error {
return pw.ws.Close()
}
func newClient(u *url.URL, method string) (*websocket.Conn, *http.Response, error) {
var c net.Conn
var err error
hostport := hostPort(u)
switch u.Scheme {
case "http":
c, err = net.Dial("tcp", hostport)
case "https":
c, err = tls.Dial("tcp", hostport, nil)
default:
return nil, nil, errors.New("unrecognized URL scheme")
}
if err != nil {
return nil, nil, err
}
return websocket.NewClient(c, u, http.Header{xVerb: []string{method}}, readBufSize, writeBufSize)
}
func hostPort(u *url.URL) string {
if strings.Contains(u.Host, ":") {
return u.Host
}
return u.Host + ":" + u.Scheme
}
// errorFromResponse returns err if err != nil, or another non-nil error if resp
// indicates a non-HTTP 200 response.
func errorFromResponse(resp *http.Response, err error) error {
if err != nil {
return err
}
if resp.StatusCode != http.StatusOK {
switch resp.StatusCode {
case http.StatusNotFound:
return os.ErrNotExist
default:
return fmt.Errorf("HTTP status %d", resp.StatusCode)
}
}
return nil
}