github.com/gramework/gramework@v1.8.1-0.20231027140105-82555c9057f5/firewall.go (about) 1 // Copyright 2017-present Kirill Danshin and Gramework contributors 2 // Copyright 2019-present Highload LTD (UK CN: 11893420) 3 // 4 // Licensed under the Apache License, Version 2.0 (the "License"); 5 // you may not use this file except in compliance with the License. 6 // You may obtain a copy of the License at 7 // 8 // http://www.apache.org/licenses/LICENSE-2.0 9 // 10 11 package gramework 12 13 import ( 14 "sync/atomic" 15 "time" 16 ) 17 18 func (app *App) initFirewall() { 19 go app.firewall.cleanupRequestsCountLoop() 20 go app.firewall.releaseBlockedRemoteAddrLoop() 21 } 22 23 // NewRequest tells the firewall, that a new request happened. 24 // True is returned, if this request should be blocked, 25 // because the IP is on the block list. 26 // The remote address is always returned for logging purpose. 27 func (fw *firewall) NewRequest(ctx *Context) (shouldBeBlocked bool, remoteAddr string) { 28 if ctx == nil || ctx.RemoteAddr().String() == "0.0.0.0" { 29 return false, "" 30 } 31 // Get the remote addresse of the request 32 remoteAddr = ctx.RemoteIP().String() 33 34 // Check if this remote address is blocked 35 if fw.isBlocked(remoteAddr) { 36 return true, remoteAddr 37 } 38 39 // Register the new request in a new goroutine 40 go fw.addRequest(remoteAddr) 41 42 return false, remoteAddr 43 } 44 45 // isBlocked checks if the remote address is blocked 46 func (fw *firewall) isBlocked(remoteAddr string) bool { 47 fw.blockListMutex.Lock() 48 49 // Check if the remote address exists in the blocked map 50 _, exists := fw.blockList[remoteAddr] 51 52 fw.blockListMutex.Unlock() 53 return exists 54 } 55 56 func (fw *firewall) addRequest(remoteAddr string) { 57 fw.requestCounterMutex.Lock() 58 defer fw.requestCounterMutex.Unlock() 59 60 count, ok := fw.requestCounter[remoteAddr] 61 if !ok { 62 count = 1 63 } 64 65 // Add the remote address to the block map, if the count 66 // reached the limit. 67 if count > atomic.LoadInt64(fw.MaxReqPerMin) { 68 // Remove the remote address from the request counter map 69 delete(fw.requestCounter, remoteAddr) 70 71 // Get the current timestamp 72 timestamp := time.Now().Unix() 73 74 // Lock the mutex 75 fw.blockListMutex.Lock() 76 defer fw.blockListMutex.Unlock() 77 78 // Add the remote address with the timestamp to the block map 79 fw.blockList[remoteAddr] = timestamp 80 81 return 82 } 83 84 // Save the incremented count to the map 85 fw.requestCounter[remoteAddr] = count + 1 86 } 87 88 func (fw *firewall) cleanupRequestsCountLoop() { 89 for { 90 time.Sleep(time.Minute) 91 92 fw.requestCounterMutex.Lock() 93 94 // Clear the map if not empty 95 if len(fw.requestCounter) > 0 { 96 fw.requestCounter = make(map[string]int64) 97 } 98 fw.requestCounterMutex.Unlock() 99 } 100 } 101 102 func (fw *firewall) releaseBlockedRemoteAddrLoop() { 103 for { 104 time.Sleep(time.Minute) 105 106 fw.blockListMutex.Lock() 107 108 if len(fw.blockList) == 0 { 109 // nothing to do here 110 fw.blockListMutex.Unlock() 111 return 112 } 113 114 releaseTimestamp := time.Now().Unix() - atomic.LoadInt64(fw.BlockTimeout) 115 116 // removing expired blocks 117 for key, timestamp := range fw.blockList { 118 if timestamp < releaseTimestamp { 119 delete(fw.blockList, key) 120 } 121 } 122 fw.blockListMutex.Unlock() 123 } 124 }