ooni-probe-cli/internal/archival/dialer.go

package archival

//
// Saves dial and net.Conn events
//

import (
	"context"
	"net"
	"time"

	"github.com/ooni/probe-cli/v3/internal/model"
	"github.com/ooni/probe-cli/v3/internal/netxlite"
)

// NetworkEvent contains a network event. This kind of events
// are generated by Dialer, QUICDialer, Conn, QUICConn.
type NetworkEvent struct {
	Count      int
	Failure    error
	Finished   time.Time
	Network    string
	Operation  string
	RemoteAddr string
	Started    time.Time
}

// DialContext dials with the given dialer with the given arguments
// and stores the dial result inside of this saver.
func (s *Saver) DialContext(ctx context.Context,
	dialer model.Dialer, network, address string) (net.Conn, error) {
	started := time.Now()
	conn, err := dialer.DialContext(ctx, network, address)
	s.appendNetworkEvent(&NetworkEvent{
		Count:      0,
		Failure:    err,
		Finished:   time.Now(),
		Network:    network,
		Operation:  netxlite.ConnectOperation,
		RemoteAddr: address,
		Started:    started,
	})
	return conn, err
}

// Read reads from the given conn and stores the results in the saver.
func (s *Saver) Read(conn net.Conn, buf []byte) (int, error) {
	network := conn.RemoteAddr().Network()
	remoteAddr := conn.RemoteAddr().String()
	started := time.Now()
	count, err := conn.Read(buf)
	s.appendNetworkEvent(&NetworkEvent{
		Count:      count,
		Failure:    err,
		Finished:   time.Now(),
		Network:    network,
		Operation:  netxlite.ReadOperation,
		RemoteAddr: remoteAddr,
		Started:    started,
	})
	return count, err
}

// Write writes to the given conn and stores the results into the saver.
func (s *Saver) Write(conn net.Conn, buf []byte) (int, error) {
	network := conn.RemoteAddr().Network()
	remoteAddr := conn.RemoteAddr().String()
	started := time.Now()
	count, err := conn.Write(buf)
	s.appendNetworkEvent(&NetworkEvent{
		Count:      count,
		Failure:    err,
		Finished:   time.Now(),
		Network:    network,
		Operation:  netxlite.WriteOperation,
		RemoteAddr: remoteAddr,
		Started:    started,
	})
	return count, err
}

func (s *Saver) appendNetworkEvent(ev *NetworkEvent) {
	s.mu.Lock()
	s.trace.Network = append(s.trace.Network, ev)
	s.mu.Unlock()
}