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  }