|
|
@@ -13,12 +13,6 @@ import (
|
|
13
|
13
|
"github.com/9seconds/mtg/wrappers"
|
|
14
|
14
|
)
|
|
15
|
15
|
|
|
16
|
|
-var (
|
|
17
|
|
- rpcCloseExtTag = [4]byte{0xa2, 0x34, 0xb6, 0x5e}
|
|
18
|
|
- rpcProxyAnsTag = [4]byte{0x0d, 0xda, 0x03, 0x44}
|
|
19
|
|
- rpcSimpleAckTag = [4]byte{0x9b, 0x40, 0xac, 0x3b}
|
|
20
|
|
-)
|
|
21
|
|
-
|
|
22
|
16
|
type ProxyRequestReadWriteCloserWithAddr struct {
|
|
23
|
17
|
wrappers.BufferedReader
|
|
24
|
18
|
|
|
|
@@ -35,38 +29,53 @@ func (p *ProxyRequestReadWriteCloserWithAddr) Read(buf []byte) (int, error) {
|
|
35
|
29
|
return errors.Annotate(err, "Cannot read RPC tag")
|
|
36
|
30
|
}
|
|
37
|
31
|
|
|
38
|
|
- if bytes.Equal(ansBuf.Bytes(), rpcCloseExtTag[:]) {
|
|
39
|
|
- return errors.New("Connection has been closed remotely")
|
|
40
|
|
- } else if bytes.Equal(ansBuf.Bytes(), rpcProxyAnsTag[:]) {
|
|
41
|
|
- if _, err := io.CopyN(ioutil.Discard, p.conn, 8+4); err != nil {
|
|
42
|
|
- return errors.Annotate(err, "Cannot skip flags and connid")
|
|
43
|
|
- }
|
|
44
|
|
- for {
|
|
45
|
|
- n, err := p.conn.Read(buf)
|
|
46
|
|
- if err != nil {
|
|
47
|
|
- return errors.Annotate(err, "Cannot read proxy answer")
|
|
48
|
|
- }
|
|
49
|
|
- if n == 0 {
|
|
50
|
|
- break
|
|
51
|
|
- }
|
|
52
|
|
- p.Buffer.Write(buf[:n])
|
|
53
|
|
- }
|
|
54
|
|
- return nil
|
|
55
|
|
- } else if bytes.Equal(ansBuf.Bytes(), rpcSimpleAckTag[:]) {
|
|
56
|
|
- if _, err := io.CopyN(ioutil.Discard, p.conn, 8); err != nil {
|
|
57
|
|
- return errors.Annotate(err, "Cannot skip connid")
|
|
58
|
|
- }
|
|
59
|
|
- if _, err := io.CopyN(p.Buffer, p.conn, 4); err != nil {
|
|
60
|
|
- return errors.Annotate(err, "Cannot read simple ack")
|
|
61
|
|
- }
|
|
62
|
|
- p.req.Options.SimpleAck = true
|
|
63
|
|
- return nil
|
|
|
32
|
+ if bytes.Equal(ansBuf.Bytes(), rpc.RPCTagCloseExt) {
|
|
|
33
|
+ return p.readCloseExt()
|
|
|
34
|
+ } else if bytes.Equal(ansBuf.Bytes(), rpc.RPCTagProxyAns) {
|
|
|
35
|
+ return p.readProxyAns(buf)
|
|
|
36
|
+ } else if bytes.Equal(ansBuf.Bytes(), rpc.RPCTagSimpleAck) {
|
|
|
37
|
+ return p.readSimpleAck()
|
|
64
|
38
|
}
|
|
65
|
39
|
|
|
66
|
40
|
return nil
|
|
67
|
41
|
})
|
|
68
|
42
|
}
|
|
69
|
43
|
|
|
|
44
|
+func (p *ProxyRequestReadWriteCloserWithAddr) readCloseExt() error {
|
|
|
45
|
+ return errors.New("Connection has been closed remotely")
|
|
|
46
|
+}
|
|
|
47
|
+
|
|
|
48
|
+func (p *ProxyRequestReadWriteCloserWithAddr) readProxyAns(buf []byte) error {
|
|
|
49
|
+ if _, err := io.CopyN(ioutil.Discard, p.conn, 8+4); err != nil {
|
|
|
50
|
+ return errors.Annotate(err, "Cannot skip flags and connid")
|
|
|
51
|
+ }
|
|
|
52
|
+
|
|
|
53
|
+ for {
|
|
|
54
|
+ n, err := p.conn.Read(buf)
|
|
|
55
|
+ if err != nil {
|
|
|
56
|
+ return errors.Annotate(err, "Cannot read proxy answer")
|
|
|
57
|
+ }
|
|
|
58
|
+ if n == 0 {
|
|
|
59
|
+ break
|
|
|
60
|
+ }
|
|
|
61
|
+ p.Buffer.Write(buf[:n])
|
|
|
62
|
+ }
|
|
|
63
|
+
|
|
|
64
|
+ return nil
|
|
|
65
|
+}
|
|
|
66
|
+
|
|
|
67
|
+func (p *ProxyRequestReadWriteCloserWithAddr) readSimpleAck() error {
|
|
|
68
|
+ if _, err := io.CopyN(ioutil.Discard, p.conn, 8); err != nil {
|
|
|
69
|
+ return errors.Annotate(err, "Cannot skip connid")
|
|
|
70
|
+ }
|
|
|
71
|
+ if _, err := io.CopyN(p.Buffer, p.conn, 4); err != nil {
|
|
|
72
|
+ return errors.Annotate(err, "Cannot read simple ack")
|
|
|
73
|
+ }
|
|
|
74
|
+ p.req.Options.SimpleAck = true
|
|
|
75
|
+
|
|
|
76
|
+ return nil
|
|
|
77
|
+}
|
|
|
78
|
+
|
|
70
|
79
|
func (p *ProxyRequestReadWriteCloserWithAddr) Write(raw []byte) (int, error) {
|
|
71
|
80
|
if _, err := p.conn.Write(p.req.Bytes(raw)); err != nil {
|
|
72
|
81
|
return 0, err
|