Golang lob.Packet類代碼示例

func (t *transport) handlePing(ch *e3x.Channel) {
	var (
		err   error
		pkt   *lob.Packet
		id    string
		alive bool
	)

	defer ch.Close()

	pkt, err = ch.ReadPacket()
	if err != nil {
		// log error
		// tracef("error: %s", err)
		return
	}

	id, _ = pkt.Header().GetString("vn")
	rpc := t.lookupRPC(id)
	if rpc == nil {
		alive = false
	} else {
		alive = true
	}

	pkt = &lob.Packet{}
	pkt.Header().SetBool("alive", alive)
	err = ch.WritePacket(pkt)
	if err != nil {
		// log error
		// tracef("error: %s", err)
		return
	}
}

func (c *Channel) traceReceivedPacket(pkt *lob.Packet) {
	if tracer.Enabled {
		tracer.Emit("channel.rcv.packet", tracer.Info{
			"channel_id": c.TID,
			"packet_id":  pkt.TID,
			"packet": tracer.Info{
				"header": pkt.Header(),
				"body":   base64.StdEncoding.EncodeToString(pkt.Body(nil)),
			},
		})
	}
}

func TestPingPong(t *testing.T) {
	// t.Parallel()
	logs.ResetLogger()

	withTwoEndpoints(t, func(A, B *Endpoint) {
		var (
			assert = assert.New(t)
			c      *Channel
			ident  *Identity
			pkt    *lob.Packet
			err    error
		)

		go func() {
			c, err := A.Listen("ping", false).AcceptChannel()

			c.SetDeadline(time.Now().Add(10 * time.Second))

			if assert.NoError(err) && assert.NotNil(c) {
				defer c.Close()

				pkt, err = c.ReadPacket()
				if assert.NoError(err) && assert.NotNil(pkt) {
					assert.Equal("ping", string(pkt.Body(nil)))

					err = c.WritePacket(lob.New([]byte("pong")))
					assert.NoError(err)
				}
			}
		}()

		ident, err = A.LocalIdentity()
		assert.NoError(err)

		c, err = B.Open(ident, "ping", false)
		assert.NoError(err)
		if assert.NotNil(c) {
			defer c.Close()

			c.SetDeadline(time.Now().Add(10 * time.Second))

			err = c.WritePacket(lob.New([]byte("ping")))
			assert.NoError(err)

			pkt, err = c.ReadPacket()
			if assert.NoError(err) && assert.NotNil(pkt) {
				assert.Equal("pong", string(pkt.Body(nil)))
			}
		}
	})
}

func (t *transport) handleGetPredecessor(ch *e3x.Channel) {
	var (
		err   error
		pkt   *lob.Packet
		id    string
		vnode *chord.Vnode
		res   *completeVnode
	)

	defer ch.Close()

	pkt, err = ch.ReadPacket()
	if err != nil {
		// log error
		// tracef("error: %s", err)
		return
	}

	id, _ = pkt.Header().GetString("vn")
	rpc := t.lookupRPC(id)
	if rpc == nil {
		// log
		// tracef("error: %s", "no RPC")
		return
	}

	vnode, err = rpc.GetPredecessor()
	if err != nil {
		// log
		// tracef("error: %s", err)
		return
	}

	res = t.completeVnode(vnode)
	err = json.NewEncoder(newStream(ch)).Encode(&res)
	if err != nil {
		// log
		// tracef("error: %s", err)
		return
	}

	if res != nil {
		// tracef("handle.GetPredecessor(Vnode(%q)) => Vnode(%q)", id, res.Id)
	}
}

func (c *Channel) traceDroppedPacket(pkt *lob.Packet, reason string) {
	if tracer.Enabled {
		info := tracer.Info{
			"channel_id": c.TID,
			"packet_id":  pkt.TID,
			"reason":     reason,
		}

		if pkt != nil {
			info["packet"] = tracer.Info{
				"header": pkt.Header(),
				"body":   base64.StdEncoding.EncodeToString(pkt.Body(nil)),
			}
		}

		tracer.Emit("channel.rcv.packet", info)
	}
}

func (c *Channel) applyAckHeaders(pkt *lob.Packet) {
	if !c.reliable {
		return
	}

	if c.iSeq == cBlankSeq {
		// nothin to ack
		return
	}

	hdr := pkt.Header()
	if c.iSeq >= cInitialSeq {
		hdr.Ack, hdr.HasAck = c.iSeq, true
	}
	if l := c.buildMissList(); len(l) > 0 {
		hdr.Miss, hdr.HasMiss = l, true
	}

	c.iAckedSeq = c.iSeq
}

// Request a nodes predecessor
func (t *transport) GetPredecessor(vn *chord.Vnode) (*chord.Vnode, error) {
	var (
		addr *e3x.Addr
		ch   *e3x.Channel
		pkt  *lob.Packet
		res  *completeVnode
		err  error
	)

	addr = t.lookupAddr(hashname.H(vn.Host))
	if addr == nil {
		return nil, e3x.ErrNoAddress
	}

	ch, err = t.e.Open(addr, "chord.predecessor.get", true)
	if err != nil {
		return nil, err
	}

	defer ch.Close()

	// ch.SetReadDeadline(time.Now().Add(30*time.Second))
	// ch.SetWriteDeadline(time.Now().Add(30*time.Second))

	pkt = &lob.Packet{}
	pkt.Header().SetString("vn", vn.String())
	err = ch.WritePacket(pkt)
	if err != nil {
		return nil, err
	}

	err = json.NewDecoder(newStream(ch)).Decode(&res)
	if err != nil {
		return nil, err
	}

	if res != nil {
		// tracef("GetPredecessor(Vnode(%q)) => Vnode(%q)", vn.String(), res.Id)
	}
	return t.internalVnode(res), nil
}

func (c *Channel) traceWrite(pkt *lob.Packet, p *Pipe) {
	if tracer.Enabled {
		info := tracer.Info{
			"channel_id": c.TID,
		}

		if p != nil {
			info["path"] = p.RemoteAddr().String()
		}

		if pkt != nil {
			info["packet_id"] = pkt.TID
			info["packet"] = tracer.Info{
				"header": pkt.Header(),
				"body":   base64.StdEncoding.EncodeToString(pkt.Body(nil)),
			}
		}

		tracer.Emit("channel.write", info)
	}
}

// Ping a Vnode, check for liveness
func (t *transport) Ping(vn *chord.Vnode) (bool, error) {
	var (
		addr  *e3x.Addr
		ch    *e3x.Channel
		pkt   *lob.Packet
		alive bool
		err   error
	)

	addr = t.lookupAddr(hashname.H(vn.Host))
	if addr == nil {
		return false, e3x.ErrNoAddress
	}

	ch, err = t.e.Open(addr, "chord.ping", true)
	if err != nil {
		return false, err
	}

	defer ch.Close()

	// ch.SetReadDeadline(time.Now().Add(30*time.Second))
	// ch.SetWriteDeadline(time.Now().Add(30*time.Second))

	pkt = &lob.Packet{}
	pkt.Header().SetString("vn", vn.String())
	err = ch.WritePacket(pkt)
	if err != nil {
		return false, err
	}

	pkt, err = ch.ReadPacket()
	if err != nil {
		return false, err
	}
	alive, _ = pkt.Header().GetBool("alive")
	// tracef("Ping(Vnode(%q)) => %v", vn.String(), alive)
	return alive, nil
}

func (c *Channel) traceWriteError(pkt *lob.Packet, p *Pipe, reason error) error {
	if tracer.Enabled {
		info := tracer.Info{
			"channel_id": c.TID,
			"reason":     reason.Error(),
		}

		if p != nil {
			info["path"] = p.RemoteAddr().String()
		}

		if pkt != nil {
			info["packet_id"] = pkt.TID
			info["packet"] = tracer.Info{
				"header": pkt.Header(),
				"body":   base64.StdEncoding.EncodeToString(pkt.Body(nil)),
			}
		}

		tracer.Emit("channel.write.error", info)
	}
	return reason
}

func (s *state) DecryptPacket(pkt *lob.Packet) (*lob.Packet, error) {
	s.mtx.RLock()
	defer s.mtx.RUnlock()

	if !s.CanDecryptPacket() {
		return nil, cipherset.ErrInvalidState
	}
	if pkt == nil {
		return nil, nil
	}

	if !pkt.Header().IsZero() || pkt.BodyLen() < lenToken+lenNonce {
		return nil, cipherset.ErrInvalidPacket
	}

	var (
		nonce    [lenNonce]byte
		bodyRaw  []byte
		innerRaw []byte
		innerPkt *lob.Packet
		body     = bufpool.New()
		inner    = bufpool.New()
		ok       bool
	)

	pkt.Body(body.SetLen(pkt.BodyLen()).RawBytes()[:0])
	bodyRaw = body.RawBytes()
	innerRaw = inner.RawBytes()

	// compare token
	if !bytes.Equal(bodyRaw[:lenToken], (*s.localToken)[:]) {
		inner.Free()
		body.Free()
		return nil, cipherset.ErrInvalidPacket
	}

	// copy nonce
	copy(nonce[:], bodyRaw[lenToken:lenToken+lenNonce])

	// decrypt inner packet
	innerRaw, ok = box.OpenAfterPrecomputation(
		innerRaw[:0], bodyRaw[lenToken+lenNonce:], &nonce, s.lineDecryptionKey)
	if !ok {
		inner.Free()
		body.Free()
		return nil, cipherset.ErrInvalidPacket
	}
	inner.SetLen(len(innerRaw))

	innerPkt, err := lob.Decode(inner)
	if err != nil {
		inner.Free()
		body.Free()
		return nil, err
	}

	inner.Free()
	body.Free()

	return innerPkt, nil
}

func (m *MockExchange) deliverPacket(pkt *lob.Packet, dst *Pipe) error {
	pkt.TID = 0
	args := m.Called(pkt)
	return args.Error(0)
}

func BenchmarkReadWriteUnreliable(b *testing.B) {
	defer dumpExpVar(b)
	logs.ResetLogger()

	withTwoEndpoints(b, func(A, B *Endpoint) {
		A.setOptions(DisableLog())
		B.setOptions(DisableLog())

		var (
			c     *Channel
			ident *Identity
			pkt   *lob.Packet
			err   error
			body  = bytes.Repeat([]byte{'x'}, 1300)
		)

		b.SetBytes(int64(len(body)))
		b.ResetTimer()

		go func() {
			c, err := A.Listen("flood", false).AcceptChannel()
			if err != nil {
				b.Fatal(err)
			}

			defer c.Close()

			pkt, err = c.ReadPacket()
			if err != nil {
				b.Fatal(err)
			}

			for i := 0; i < b.N; i++ {
				pkt := lob.New(body)
				err = c.WritePacket(pkt)
				if err != nil {
					b.Fatal(err)
				}

				// Give the other go routines some room to breath when GOMAXPROCS=1
				runtime.Gosched()
			}
		}()

		ident, err = A.LocalIdentity()
		if err != nil {
			b.Fatal(err)
		}

		c, err = B.Open(ident, "flood", false)
		if err != nil {
			b.Fatal(err)
		}

		defer c.Close()

		err = c.WritePacket(lob.New(nil))
		if err != nil {
			b.Fatal(err)
		}

		for {
			pkt, err = c.ReadPacket()
			if err == io.EOF {
				break
			}
			if err != nil {
				b.Fatal(err)
			}
			pkt.Free()
		}

		b.StopTimer()
	})
}

func (s *state) DecryptPacket(pkt *lob.Packet) (*lob.Packet, error) {
	s.mtx.RLock()
	defer s.mtx.RUnlock()

	if !s.CanDecryptPacket() {
		return nil, cipherset.ErrInvalidState
	}
	if pkt == nil {
		return nil, nil
	}

	if !pkt.Header().IsZero() || pkt.BodyLen() < 16+4+4 {
		return nil, cipherset.ErrInvalidPacket
	}

	var (
		nonce    [16]byte
		bodyRaw  []byte
		innerRaw []byte
		innerLen = pkt.BodyLen() - (16 + 4 + 4)
		body     = bufpool.New()
		inner    = bufpool.New().SetLen(innerLen)
	)

	pkt.Body(body.SetLen(pkt.BodyLen()).RawBytes()[:0])
	bodyRaw = body.RawBytes()
	innerRaw = inner.RawBytes()

	// compare token
	if !bytes.Equal(bodyRaw[:16], (*s.localToken)[:]) {
		inner.Free()
		body.Free()
		return nil, cipherset.ErrInvalidPacket
	}

	// copy nonce
	copy(nonce[:], bodyRaw[16:16+4])

	{ // verify hmac
		mac := bodyRaw[16+4+innerLen:]

		macKey := append(s.lineDecryptionKey, nonce[:4]...)

		h := hmac.New(sha256.New, macKey)
		h.Write(bodyRaw[16+4 : 16+4+innerLen])
		if subtle.ConstantTimeCompare(mac, fold(h.Sum(nil), 4)) != 1 {
			inner.Free()
			body.Free()
			return nil, cipherset.ErrInvalidPacket
		}
	}

	{ // decrypt inner
		aesBlock, err := aes.NewCipher(s.lineDecryptionKey)
		if err != nil {
			inner.Free()
			body.Free()
			return nil, err
		}

		aes := Cipher.NewCTR(aesBlock, nonce[:])
		if aes == nil {
			inner.Free()
			body.Free()
			return nil, cipherset.ErrInvalidPacket
		}

		aes.XORKeyStream(innerRaw, bodyRaw[16+4:16+4+innerLen])
	}

	innerPkt, err := lob.Decode(inner)
	if err != nil {
		inner.Free()
		body.Free()
		return nil, err
	}

	inner.Free()
	body.Free()

	return innerPkt, nil
}

func (c *Channel) write(pkt *lob.Packet, p *Pipe) error {
	if pkt.TID == 0 {
		pkt.TID = tracer.NewID()
	}

	if c.broken {
		// When a channel is marked as broken the all writes
		// must return a BrokenChannelError.
		return c.traceWriteError(pkt, p,
			&BrokenChannelError{c.hashname, c.typ, c.id})
	}

	if c.writeDeadlineReached {
		// When a channel reached a write deadline then all writes
		// must return a ErrTimeout.
		return c.traceWriteError(pkt, p,
			ErrTimeout)
	}

	if c.deliveredEnd {
		// When a channel sent a packet with the "end" header set
		// then all subsequent writes must return io.EOF
		return c.traceWriteError(pkt, p,
			io.EOF)
	}

	c.oSeq++
	hdr := pkt.Header()
	hdr.C, hdr.HasC = c.id, true
	if c.reliable {
		hdr.Seq, hdr.HasSeq = c.oSeq, true
	}
	if !c.serverside && c.oSeq == cInitialSeq {
		hdr.Type, hdr.HasType = c.typ, true
	}

	end := hdr.HasEnd && hdr.End
	if end {
		c.deliveredEnd = true
		c.setCloseDeadline()
	}

	if c.reliable {
		if c.oSeq%30 == 0 || hdr.End {
			c.applyAckHeaders(pkt)
		}
		c.writeBuffer[c.oSeq] = &writeBufferEntry{pkt, end, time.Time{}, p}
		c.needsResend = false
	}

	err := c.x.deliverPacket(pkt, p)
	if err != nil {
		return c.traceWriteError(pkt, p, err)
	}
	statChannelSndPkt.Add(1)
	if pkt.Header().HasAck {
		statChannelSndAckInline.Add(1)
	}

	if c.oSeq == cInitialSeq && c.serverside {
		c.unsetOpenDeadline()
	}

	c.traceWrite(pkt, p)

	if !c.reliable {
		pkt.Free()
	}

	return nil
}