extras/hs-test/address_allocator.go

   1 package main
   2
   3 import (
   4         "errors"
   5         "fmt"
   6         "os"
   7         "strings"
   8
   9         "github.com/edwarnicke/exechelper"
  10 )
  11
  12 type AddressCounter = int
  13
  14 type Ip4AddressAllocator struct {
  15         networks    map[int]AddressCounter
  16         chosenOctet int
  17         assignedIps []string
  18 }
  19
  20 func (a *Ip4AddressAllocator) AddNetwork(networkNumber int) {
  21         a.networks[networkNumber] = 1
  22 }
  23
  24 func (a *Ip4AddressAllocator) NewIp4InterfaceAddress(inputNetworkNumber ...int) (string, error) {
  25         var networkNumber int = 0
  26         if len(inputNetworkNumber) > 0 {
  27                 networkNumber = inputNetworkNumber[0]
  28         }
  29
  30         if _, ok := a.networks[networkNumber]; !ok {
  31                 a.AddNetwork(networkNumber)
  32         }
  33
  34         numberOfAddresses := a.networks[networkNumber]
  35
  36         if numberOfAddresses == 254 {
  37                 return "", fmt.Errorf("no available IPv4 addresses")
  38         }
  39
  40         address, err := a.createIpAddress(networkNumber, numberOfAddresses)
  41
  42         a.networks[networkNumber] = numberOfAddresses + 1
  43
  44         return address + "/24", err
  45 }
  46
  47 // Creates a file every time an IP is assigned: used to keep track of addresses in use.
  48 // If an address is not in use, 'counter' is then copied to 'chosenOctet' and it is used for the remaining tests.
  49 // Also checks host IP addresses.
  50 func (a *Ip4AddressAllocator) createIpAddress(networkNumber int, numberOfAddresses int) (string, error) {
  51         hostIps, _ := exechelper.CombinedOutput("ip a")
  52         counter := 10
  53         var address string
  54
  55         for {
  56                 if a.chosenOctet != 0 {
  57                         address = fmt.Sprintf("10.%v.%v.%v", a.chosenOctet, networkNumber, numberOfAddresses)
  58                         file, err := os.Create(address)
  59                         if err != nil {
  60                                 return "", errors.New("unable to create file: " + fmt.Sprint(err))
  61                         }
  62                         file.Close()
  63                         break
  64                 } else {
  65                         address = fmt.Sprintf("10.%v.%v.%v", counter, networkNumber, numberOfAddresses)
  66                         _, err := os.Stat(address)
  67                         if err == nil || strings.Contains(string(hostIps), address) {
  68                                 counter++
  69                         } else if os.IsNotExist(err) {
  70                                 file, err := os.Create(address)
  71                                 if err != nil {
  72                                         return "", errors.New("unable to create file: " + fmt.Sprint(err))
  73                                 }
  74                                 file.Close()
  75                                 a.chosenOctet = counter
  76                                 break
  77                         } else {
  78                                 return "", errors.New("an error occurred while checking if a file exists: " + fmt.Sprint(err))
  79                         }
  80                 }
  81         }
  82
  83         a.assignedIps = append(a.assignedIps, address)
  84         return address, nil
  85 }
  86
  87 func (a *Ip4AddressAllocator) deleteIpAddresses() {
  88         for ip := range a.assignedIps {
  89                 os.Remove(a.assignedIps[ip])
  90         }
  91 }
  92
  93 func NewIp4AddressAllocator() *Ip4AddressAllocator {
  94         var ip4AddrAllocator = new(Ip4AddressAllocator)
  95         ip4AddrAllocator.networks = make(map[int]AddressCounter)
  96         ip4AddrAllocator.AddNetwork(0)
  97         return ip4AddrAllocator
  98 }