forked from hyper0x/goc2p
-
Notifications
You must be signed in to change notification settings - Fork 0
/
comm.go
132 lines (123 loc) · 3.11 KB
/
comm.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
package testhelper
import (
"bufio"
"bytes"
"encoding/json"
"fmt"
loadgenlib "loadgen/lib"
"math/rand"
"net"
"time"
)
const (
DELIM = '\n'
)
type TcpComm struct {
addr string
}
func NewTcpComm(addr string) loadgenlib.Caller {
return &TcpComm{addr: addr}
}
func (comm *TcpComm) BuildReq() loadgenlib.RawReq {
id := time.Now().UnixNano()
sreq := ServerReq{
Id: id,
Operands: []int{
int(rand.Int31n(1000)),
int(rand.Int31n(1000))},
Operator: func() string {
op := []string{"+", "-", "*", "/"}
return op[rand.Int31n(100)%4]
}(),
}
bytes, err := json.Marshal(sreq)
if err != nil {
panic(err)
}
rawReq := loadgenlib.RawReq{Id: id, Req: bytes}
return rawReq
}
func (comm *TcpComm) Call(req []byte, timeoutNs time.Duration) ([]byte, error) {
conn, err := net.DialTimeout("tcp", comm.addr, timeoutNs)
if err != nil {
return nil, err
}
_, err = write(conn, req, DELIM)
if err != nil {
return nil, err
}
return read(conn, DELIM)
}
func (comm *TcpComm) CheckResp(
rawReq loadgenlib.RawReq, rawResp loadgenlib.RawResp) *loadgenlib.CallResult {
var commResult loadgenlib.CallResult
commResult.Id = rawResp.Id
commResult.Req = rawReq
commResult.Resp = rawResp
var sreq ServerReq
err := json.Unmarshal(rawReq.Req, &sreq)
if err != nil {
commResult.Code = loadgenlib.RESULT_CODE_FATAL_CALL
commResult.Msg =
fmt.Sprintf("Incorrectly formatted Req: %s!\n", string(rawReq.Req))
return &commResult
}
var sresp ServerResp
err = json.Unmarshal(rawResp.Resp, &sresp)
if err != nil {
commResult.Code = loadgenlib.RESULT_CODE_ERROR_RESPONSE
commResult.Msg =
fmt.Sprintf("Incorrectly formatted Resp: %s!\n", string(rawResp.Resp))
return &commResult
}
if sresp.Id != sreq.Id {
commResult.Code = loadgenlib.RESULT_CODE_ERROR_RESPONSE
commResult.Msg =
fmt.Sprintf("Inconsistent raw id! (%d != %d)\n", rawReq.Id, rawResp.Id)
return &commResult
}
if sresp.Err != nil {
commResult.Code = loadgenlib.RESULT_CODE_ERROR_CALEE
commResult.Msg =
fmt.Sprintf("Abnormal server: %s!\n", sresp.Err)
return &commResult
}
if sresp.Result != op(sreq.Operands, sreq.Operator) {
commResult.Code = loadgenlib.RESULT_CODE_ERROR_RESPONSE
commResult.Msg =
fmt.Sprintf(
"Incorrect result: %s!\n",
genFormula(sreq.Operands, sreq.Operator, sresp.Result, false))
return &commResult
}
commResult.Code = loadgenlib.RESULT_CODE_SUCCESS
commResult.Msg = fmt.Sprintf("Success. (%s)", sresp.Formula)
return &commResult
}
func read(conn net.Conn, delim byte) ([]byte, error) {
readBytes := make([]byte, 1)
var buffer bytes.Buffer
for {
_, err := conn.Read(readBytes)
if err != nil {
return nil, err
}
readByte := readBytes[0]
if readByte == delim {
break
}
buffer.WriteByte(readByte)
}
return buffer.Bytes(), nil
}
func write(conn net.Conn, content []byte, delim byte) (int, error) {
writer := bufio.NewWriter(conn)
n, err := writer.Write(content)
if err == nil {
writer.WriteByte(delim)
}
if err == nil {
err = writer.Flush()
}
return n, err
}