Created
March 2, 2017 08:16
-
-
Save eviltofu/c3441eef2ea58d16725f7ef0b2ba1adf to your computer and use it in GitHub Desktop.
Simple Connection Rate Limiter
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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")) | |
} |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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 | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment