github.com/ylsGit/go-ethereum@v1.6.5/rpc/subscription_test.go (about) 1 // Copyright 2016 The go-ethereum Authors 2 // This file is part of the go-ethereum library. 3 // 4 // The go-ethereum library is free software: you can redistribute it and/or modify 5 // it under the terms of the GNU Lesser General Public License as published by 6 // the Free Software Foundation, either version 3 of the License, or 7 // (at your option) any later version. 8 // 9 // The go-ethereum library is distributed in the hope that it will be useful, 10 // but WITHOUT ANY WARRANTY; without even the implied warranty of 11 // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 12 // GNU Lesser General Public License for more details. 13 // 14 // You should have received a copy of the GNU Lesser General Public License 15 // along with the go-ethereum library. If not, see <http://www.gnu.org/licenses/>. 16 17 package rpc 18 19 import ( 20 "context" 21 "encoding/json" 22 "fmt" 23 "net" 24 "sync" 25 "testing" 26 "time" 27 ) 28 29 type NotificationTestService struct { 30 mu sync.Mutex 31 unsubscribed bool 32 33 gotHangSubscriptionReq chan struct{} 34 unblockHangSubscription chan struct{} 35 } 36 37 func (s *NotificationTestService) Echo(i int) int { 38 return i 39 } 40 41 func (s *NotificationTestService) wasUnsubCallbackCalled() bool { 42 s.mu.Lock() 43 defer s.mu.Unlock() 44 return s.unsubscribed 45 } 46 47 func (s *NotificationTestService) Unsubscribe(subid string) { 48 s.mu.Lock() 49 s.unsubscribed = true 50 s.mu.Unlock() 51 } 52 53 func (s *NotificationTestService) SomeSubscription(ctx context.Context, n, val int) (*Subscription, error) { 54 notifier, supported := NotifierFromContext(ctx) 55 if !supported { 56 return nil, ErrNotificationsUnsupported 57 } 58 59 // by explicitly creating an subscription we make sure that the subscription id is send back to the client 60 // before the first subscription.Notify is called. Otherwise the events might be send before the response 61 // for the eth_subscribe method. 62 subscription := notifier.CreateSubscription() 63 64 go func() { 65 // test expects n events, if we begin sending event immediately some events 66 // will probably be dropped since the subscription ID might not be send to 67 // the client. 68 time.Sleep(5 * time.Second) 69 for i := 0; i < n; i++ { 70 if err := notifier.Notify(subscription.ID, val+i); err != nil { 71 return 72 } 73 } 74 75 select { 76 case <-notifier.Closed(): 77 s.mu.Lock() 78 s.unsubscribed = true 79 s.mu.Unlock() 80 case <-subscription.Err(): 81 s.mu.Lock() 82 s.unsubscribed = true 83 s.mu.Unlock() 84 } 85 }() 86 87 return subscription, nil 88 } 89 90 // HangSubscription blocks on s.unblockHangSubscription before 91 // sending anything. 92 func (s *NotificationTestService) HangSubscription(ctx context.Context, val int) (*Subscription, error) { 93 notifier, supported := NotifierFromContext(ctx) 94 if !supported { 95 return nil, ErrNotificationsUnsupported 96 } 97 98 s.gotHangSubscriptionReq <- struct{}{} 99 <-s.unblockHangSubscription 100 subscription := notifier.CreateSubscription() 101 102 go func() { 103 notifier.Notify(subscription.ID, val) 104 }() 105 return subscription, nil 106 } 107 108 func TestNotifications(t *testing.T) { 109 server := NewServer() 110 service := &NotificationTestService{} 111 112 if err := server.RegisterName("eth", service); err != nil { 113 t.Fatalf("unable to register test service %v", err) 114 } 115 116 clientConn, serverConn := net.Pipe() 117 118 go server.ServeCodec(NewJSONCodec(serverConn), OptionMethodInvocation|OptionSubscriptions) 119 120 out := json.NewEncoder(clientConn) 121 in := json.NewDecoder(clientConn) 122 123 n := 5 124 val := 12345 125 request := map[string]interface{}{ 126 "id": 1, 127 "method": "eth_subscribe", 128 "version": "2.0", 129 "params": []interface{}{"someSubscription", n, val}, 130 } 131 132 // create subscription 133 if err := out.Encode(request); err != nil { 134 t.Fatal(err) 135 } 136 137 var subid string 138 response := jsonSuccessResponse{Result: subid} 139 if err := in.Decode(&response); err != nil { 140 t.Fatal(err) 141 } 142 143 var ok bool 144 if _, ok = response.Result.(string); !ok { 145 t.Fatalf("expected subscription id, got %T", response.Result) 146 } 147 148 for i := 0; i < n; i++ { 149 var notification jsonNotification 150 if err := in.Decode(¬ification); err != nil { 151 t.Fatalf("%v", err) 152 } 153 154 if int(notification.Params.Result.(float64)) != val+i { 155 t.Fatalf("expected %d, got %d", val+i, notification.Params.Result) 156 } 157 } 158 159 clientConn.Close() // causes notification unsubscribe callback to be called 160 time.Sleep(1 * time.Second) 161 162 if !service.wasUnsubCallbackCalled() { 163 t.Error("unsubscribe callback not called after closing connection") 164 } 165 } 166 167 func waitForMessages(t *testing.T, in *json.Decoder, successes chan<- jsonSuccessResponse, 168 failures chan<- jsonErrResponse, notifications chan<- jsonNotification) { 169 170 // read and parse server messages 171 for { 172 var rmsg json.RawMessage 173 if err := in.Decode(&rmsg); err != nil { 174 return 175 } 176 177 var responses []map[string]interface{} 178 if rmsg[0] == '[' { 179 if err := json.Unmarshal(rmsg, &responses); err != nil { 180 t.Fatalf("Received invalid message: %s", rmsg) 181 } 182 } else { 183 var msg map[string]interface{} 184 if err := json.Unmarshal(rmsg, &msg); err != nil { 185 t.Fatalf("Received invalid message: %s", rmsg) 186 } 187 responses = append(responses, msg) 188 } 189 190 for _, msg := range responses { 191 // determine what kind of msg was received and broadcast 192 // it to over the corresponding channel 193 if _, found := msg["result"]; found { 194 successes <- jsonSuccessResponse{ 195 Version: msg["jsonrpc"].(string), 196 Id: msg["id"], 197 Result: msg["result"], 198 } 199 continue 200 } 201 if _, found := msg["error"]; found { 202 params := msg["params"].(map[string]interface{}) 203 failures <- jsonErrResponse{ 204 Version: msg["jsonrpc"].(string), 205 Id: msg["id"], 206 Error: jsonError{int(params["subscription"].(float64)), params["message"].(string), params["data"]}, 207 } 208 continue 209 } 210 if _, found := msg["params"]; found { 211 params := msg["params"].(map[string]interface{}) 212 notifications <- jsonNotification{ 213 Version: msg["jsonrpc"].(string), 214 Method: msg["method"].(string), 215 Params: jsonSubscription{params["subscription"].(string), params["result"]}, 216 } 217 continue 218 } 219 t.Fatalf("Received invalid message: %s", msg) 220 } 221 } 222 } 223 224 // TestSubscriptionMultipleNamespaces ensures that subscriptions can exists 225 // for multiple different namespaces. 226 func TestSubscriptionMultipleNamespaces(t *testing.T) { 227 var ( 228 namespaces = []string{"eth", "shh", "bzz"} 229 server = NewServer() 230 service = NotificationTestService{} 231 clientConn, serverConn = net.Pipe() 232 233 out = json.NewEncoder(clientConn) 234 in = json.NewDecoder(clientConn) 235 successes = make(chan jsonSuccessResponse) 236 failures = make(chan jsonErrResponse) 237 notifications = make(chan jsonNotification) 238 ) 239 240 // setup and start server 241 for _, namespace := range namespaces { 242 if err := server.RegisterName(namespace, &service); err != nil { 243 t.Fatalf("unable to register test service %v", err) 244 } 245 } 246 247 go server.ServeCodec(NewJSONCodec(serverConn), OptionMethodInvocation|OptionSubscriptions) 248 defer server.Stop() 249 250 // wait for message and write them to the given channels 251 go waitForMessages(t, in, successes, failures, notifications) 252 253 // create subscriptions one by one 254 n := 3 255 for i, namespace := range namespaces { 256 request := map[string]interface{}{ 257 "id": i, 258 "method": fmt.Sprintf("%s_subscribe", namespace), 259 "version": "2.0", 260 "params": []interface{}{"someSubscription", n, i}, 261 } 262 263 if err := out.Encode(&request); err != nil { 264 t.Fatalf("Could not create subscription: %v", err) 265 } 266 } 267 268 // create all subscriptions in 1 batch 269 var requests []interface{} 270 for i, namespace := range namespaces { 271 requests = append(requests, map[string]interface{}{ 272 "id": i, 273 "method": fmt.Sprintf("%s_subscribe", namespace), 274 "version": "2.0", 275 "params": []interface{}{"someSubscription", n, i}, 276 }) 277 } 278 279 if err := out.Encode(&requests); err != nil { 280 t.Fatalf("Could not create subscription in batch form: %v", err) 281 } 282 283 timeout := time.After(30 * time.Second) 284 subids := make(map[string]string, 2*len(namespaces)) 285 count := make(map[string]int, 2*len(namespaces)) 286 287 for { 288 done := true 289 for id, _ := range count { 290 if count, found := count[id]; !found || count < (2*n) { 291 done = false 292 } 293 } 294 295 if done && len(count) == len(namespaces) { 296 break 297 } 298 299 select { 300 case suc := <-successes: // subscription created 301 subids[namespaces[int(suc.Id.(float64))]] = suc.Result.(string) 302 case failure := <-failures: 303 t.Errorf("received error: %v", failure.Error) 304 case notification := <-notifications: 305 if cnt, found := count[notification.Params.Subscription]; found { 306 count[notification.Params.Subscription] = cnt + 1 307 } else { 308 count[notification.Params.Subscription] = 1 309 } 310 case <-timeout: 311 for _, namespace := range namespaces { 312 subid, found := subids[namespace] 313 if !found { 314 t.Errorf("Subscription for '%s' not created", namespace) 315 continue 316 } 317 if count, found := count[subid]; !found || count < n { 318 t.Errorf("Didn't receive all notifications (%d<%d) in time for namespace '%s'", count, n, namespace) 319 } 320 } 321 return 322 } 323 } 324 }