waku_custom_res_mgr.go

type LastTs mutex[[]int64]

REQUEST_PER_PEER = 10

func (l LastTs) check() error {
	l.lock()
	l.unlock()
	fiveMinAgo:=time.Now().Sub(5*Second).Unix()
	streamBeforeFiveMin = sort.Search(len(l), func(i int) bool { return l[i] >= fiveMinAgo })
	if len(l) - streamBeforeFiveMin > REQUEST_PER_PEER {
		return fmt.Errorf("rate limit reached")
	}
	return nil
}
func (l LastTs) Done() { // will be called on stream closure.
	l.lock()
	

	l = append(l, time.Now().Unix())
	fiveMinAgo:=time.Now().Sub(5*Second).Unix()
	streamBeforeFiveMin = sort.Search(len(l), func(i int) bool { return l[i] >= fiveMinAgo })
	l.streamInfo = l.streamInfo[streamBeforeFiveMin:]
	
	l.unlock()
	l.libp2p.stream.Done()
}

type streamLocal struct {
	libp2p.stream
	streamInfo *LastTs
}

func (stream) Done() {
	stream.libp2p.stream.Done()
	stream.streamInfo.Done(time.Now())
} 

type WakuResMgr struct {
	ResourceManager
	wakuMu sync.Mutex

	rateLimits map[peer.ID]LastTs
}

func (mgr WakuResMgr) Openstream() (network.StreamManagementScope, error) {
	mgr.wakuMu.lock()
	if mgr.streamInfo[peer.ID] == nil {
		mgr.streamInfo[peer.ID] = new(LastTs[])
	}
	mgr.wakuMu.unlock()

	if err := mgr.streamInfo[peer.ID].check(time.Now()); err !=nil {
		return nil, err
	}
	result, err := mgr.ResourceManager.Openstream()
	if err != nil {
		return result, err
	}
	return streamLocal{result, mgr.streamInfo[peer.ID]}, nil
}