nim-libp2p/libp2p/protocols/ping.nim at master · vacp2p/nim-libp2p · GitHub

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
# SPDX-License-Identifier: Apache-2.0 OR MIT
# Copyright (c) Status Research & Development GmbH

## `Ping <https://docs.libp2p.io/concepts/protocols/#ping>`_ protocol implementation

{.push raises: [].}

import chronos, chronicles
import
  ../stream/connection, ../peerid, ../crypto/crypto, ../protocols/protocol, ../errors

export chronicles, rng, connection

logScope:
  topics = "libp2p ping"

const
  PingCodec* = "/ipfs/ping/1.0.0"
  PingSize = 32

type
  PingError* = object of LPError
  WrongPingAckError* = object of PingError

  PingHandler* = proc(peer: PeerId): Future[void] {.async: (raises: []), gcsafe.}

  Ping* = ref object of LPProtocol
    pingHandler*: PingHandler
    rng: Rng

proc new*(T: typedesc[Ping], handler: PingHandler = nil, rng: Rng): T =
  doAssert not rng.isNil, "Rng is nil"
  let ping = Ping(pinghandler: handler, rng: rng)
  ping.init()
  ping

method init*(p: Ping) =
  proc handle(stream: Stream, proto: string) {.async: (raises: [CancelledError]).} =
    try:
      trace "handling ping", stream
      var buf: array[PingSize, byte]
      while true:
        await stream.readExactly(addr buf[0], PingSize)
        trace "echoing ping", stream, pingData = @buf
        await stream.write(@buf)
        if not isNil(p.pingHandler):
          await p.pingHandler(stream.peerId)
    except LPStreamEOFError as exc:
      trace "ping stream closed", description = exc.msg, stream
    except LPStreamError as exc:
      trace "exception in ping handler", description = exc.msg, stream

  p.handler = handle
  p.codec = PingCodec

proc ping*(
    p: Ping, stream: Stream
): Future[Duration] {.
    async: (raises: [CancelledError, LPStreamError, WrongPingAckError])
.} =
  ## Sends ping to `stream`, returns the delay

  trace "initiating ping", stream

  var
    randomBuf: array[PingSize, byte]
    resultBuf: array[PingSize, byte]

  p.rng.generate(randomBuf)

  let startTime = Moment.now()

  trace "sending ping", stream
  await stream.write(@randomBuf)

  await stream.readExactly(addr resultBuf[0], PingSize)

  let responseDur = Moment.now() - startTime

  trace "got ping response", stream, responseDur

  for i in 0 ..< randomBuf.len:
    if randomBuf[i] != resultBuf[i]:
      raise newException(WrongPingAckError, "Incorrect ping data from peer!")

  trace "valid ping response", stream
  return responseDur