Simple Connection Rate Limiter

main.go

package main

import (
	"fmt"
	"net"
	"rateLimit"
	"time"
)

func main() {
	r := rateLimit.NewLimiter(1, 1*time.Minute)
	ln, err := net.Listen("tcp", ":8182")
	if err != nil {
		// handle error
		panic(err)
	}
	for {
		conn, err := ln.Accept()
		fmt.Println("Connected")
		if err != nil {
			// handle error
			panic(err)
		}
		if r.AddConnectionWithLimit(&conn) {
			go handleConnection(conn)
		} else {
			go closeConnection(conn)
		}
	}
}

func handleConnection(c net.Conn) {
	defer c.Close()
	c.Write([]byte("Thank you for connecting\n"))
	time.Sleep(1 * time.Minute)
	c.Write([]byte("Closing\n"))
}

func closeConnection(c net.Conn) {
	defer c.Close()
	c.Write([]byte("Too many connections\n"))
}

rateLimit.go

package rateLimit

import (
	"fmt"
	"net"
	"sync"
	"time"
)

// Limiter structure that contains all the calls
type Limiter struct {
	Connections []Connection
	Limit       int
	TimeFrame   time.Duration
	lock        *sync.Mutex
}

// Filter function that selects only matching Connection
type Filter func(Connection) bool

func withinDuration(d time.Duration) Filter {
	t1 := time.Now()
	fmt.Printf("Now is %v\n", t1)
	t2 := t1.Add(-d)
	fmt.Printf("Cut off time is %v\n", t2)
	f := func(c Connection) bool {
		return c.TimeStamp.After(t2)
	}
	return f
}

// AddConnectionWithLimit function that adds a connection to the list and returns if it
// is within the rate limit
func (l *Limiter) AddConnectionWithLimit(c *net.Conn) bool {
	l.lock.Lock()
	con := Connection{Conn: c, TimeStamp: time.Now()}
	connections := len(l.Connections)
	l.Connections = append(l.Connections, con)
	fmt.Printf("Started with %d connections\n", connections)
	l.FilterConnections(withinDuration(l.TimeFrame))
	connections = len(l.Connections)
	result := connections <= l.Limit
	fmt.Printf("Ended with %d connections\n", connections)
	l.lock.Unlock()
	return result
}

// FilterConnections runs on the filter on Connection
func (l *Limiter) FilterConnections(f Filter) {
	buff := make([]Connection, 0, len(l.Connections))
	for _, value := range l.Connections {
		if f(value) {
			buff = append(buff, value)
		}
	}
	l.Connections = buff
}

// Connection struct of a connection
type Connection struct {
	Conn      *net.Conn
	TimeStamp time.Time
}

// NewLimiter function which returns an initialized Limiter
func NewLimiter(size int, timeFrame time.Duration) *Limiter {
	l := Limiter{Limit: size, TimeFrame: timeFrame}
	c := make([]Connection, 0, size*2)
	l.Connections = c
	l.lock = &sync.Mutex{}
	return &l
}