Skip to content

Commit c0a528e

Browse files
authored
Replace fragmentRPC with appendOrMergeRPC (#557)
This will allow us to add more logic around when we split/merge messages. It will also allow us to build the outgoing rpcs as we go rather than building one giant rpc and then splitting it.
1 parent 048a4d3 commit c0a528e

File tree

3 files changed

+322
-107
lines changed

3 files changed

+322
-107
lines changed

fuzz_helpers_test.go

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
package pubsub
2+
3+
import (
4+
"encoding/binary"
5+
6+
pb "github.com/libp2p/go-libp2p-pubsub/pb"
7+
)
8+
9+
func generateU16(data *[]byte) uint16 {
10+
if len(*data) < 2 {
11+
return 0
12+
}
13+
14+
out := binary.LittleEndian.Uint16((*data)[:2])
15+
*data = (*data)[2:]
16+
return out
17+
}
18+
19+
func generateBool(data *[]byte) bool {
20+
if len(*data) < 1 {
21+
return false
22+
}
23+
24+
out := (*data)[0]&1 == 1
25+
*data = (*data)[1:]
26+
return out
27+
}
28+
29+
func generateMessage(data []byte, limit int) *pb.Message {
30+
msgSize := int(generateU16(&data)) % limit
31+
return &pb.Message{Data: make([]byte, msgSize)}
32+
}
33+
34+
func generateSub(data []byte, limit int) *pb.RPC_SubOpts {
35+
topicIDSize := int(generateU16(&data)) % limit
36+
subscribe := generateBool(&data)
37+
38+
str := string(make([]byte, topicIDSize))
39+
return &pb.RPC_SubOpts{Subscribe: &subscribe, Topicid: &str}
40+
}
41+
42+
func generateControl(data []byte, limit int) *pb.ControlMessage {
43+
numIWANTMsgs := int(generateU16(&data)) % (limit / 2)
44+
numIHAVEMsgs := int(generateU16(&data)) % (limit / 2)
45+
46+
ctl := &pb.ControlMessage{}
47+
48+
ctl.Iwant = make([]*pb.ControlIWant, 0, numIWANTMsgs)
49+
for i := 0; i < numIWANTMsgs; i++ {
50+
msgSize := int(generateU16(&data)) % limit
51+
msgCount := int(generateU16(&data)) % limit
52+
ctl.Iwant = append(ctl.Iwant, &pb.ControlIWant{})
53+
ctl.Iwant[i].MessageIDs = make([]string, 0, msgCount)
54+
for j := 0; j < msgCount; j++ {
55+
ctl.Iwant[i].MessageIDs = append(ctl.Iwant[i].MessageIDs, string(make([]byte, msgSize)))
56+
}
57+
}
58+
if ctl.Size() > limit {
59+
return &pb.ControlMessage{}
60+
}
61+
62+
ctl.Ihave = make([]*pb.ControlIHave, 0, numIHAVEMsgs)
63+
for i := 0; i < numIHAVEMsgs; i++ {
64+
msgSize := int(generateU16(&data)) % limit
65+
msgCount := int(generateU16(&data)) % limit
66+
topicSize := int(generateU16(&data)) % limit
67+
topic := string(make([]byte, topicSize))
68+
ctl.Ihave = append(ctl.Ihave, &pb.ControlIHave{TopicID: &topic})
69+
70+
ctl.Ihave[i].MessageIDs = make([]string, 0, msgCount)
71+
for j := 0; j < msgCount; j++ {
72+
ctl.Ihave[i].MessageIDs = append(ctl.Ihave[i].MessageIDs, string(make([]byte, msgSize)))
73+
}
74+
}
75+
if ctl.Size() > limit {
76+
return &pb.ControlMessage{}
77+
}
78+
79+
numGraft := int(generateU16(&data)) % limit
80+
ctl.Graft = make([]*pb.ControlGraft, 0, numGraft)
81+
for i := 0; i < numGraft; i++ {
82+
topicSize := int(generateU16(&data)) % limit
83+
topic := string(make([]byte, topicSize))
84+
ctl.Graft = append(ctl.Graft, &pb.ControlGraft{TopicID: &topic})
85+
}
86+
if ctl.Size() > limit {
87+
return &pb.ControlMessage{}
88+
}
89+
90+
numPrune := int(generateU16(&data)) % limit
91+
ctl.Prune = make([]*pb.ControlPrune, 0, numPrune)
92+
for i := 0; i < numPrune; i++ {
93+
topicSize := int(generateU16(&data)) % limit
94+
topic := string(make([]byte, topicSize))
95+
ctl.Prune = append(ctl.Prune, &pb.ControlPrune{TopicID: &topic})
96+
}
97+
if ctl.Size() > limit {
98+
return &pb.ControlMessage{}
99+
}
100+
101+
return ctl
102+
}
103+
104+
func generateRPC(data []byte, limit int) *RPC {
105+
rpc := &RPC{RPC: pb.RPC{}}
106+
sizeTester := RPC{RPC: pb.RPC{}}
107+
108+
msgCount := int(generateU16(&data)) % (limit / 2)
109+
rpc.Publish = make([]*pb.Message, 0, msgCount)
110+
for i := 0; i < msgCount; i++ {
111+
msg := generateMessage(data, limit)
112+
113+
sizeTester.Publish = []*pb.Message{msg}
114+
size := sizeTester.Size()
115+
sizeTester.Publish = nil
116+
if size > limit {
117+
continue
118+
}
119+
120+
rpc.Publish = append(rpc.Publish, msg)
121+
}
122+
123+
subCount := int(generateU16(&data)) % (limit / 2)
124+
rpc.Subscriptions = make([]*pb.RPC_SubOpts, 0, subCount)
125+
for i := 0; i < subCount; i++ {
126+
sub := generateSub(data, limit)
127+
128+
sizeTester.Subscriptions = []*pb.RPC_SubOpts{sub}
129+
size := sizeTester.Size()
130+
sizeTester.Subscriptions = nil
131+
if size > limit {
132+
continue
133+
}
134+
135+
rpc.Subscriptions = append(rpc.Subscriptions, sub)
136+
}
137+
138+
ctl := generateControl(data, limit)
139+
140+
sizeTester.Control = ctl
141+
size := sizeTester.Size()
142+
sizeTester.Control = nil
143+
if size <= limit {
144+
rpc.Control = ctl
145+
146+
}
147+
148+
return rpc
149+
}

0 commit comments

Comments
 (0)