init

pkg/ipset/ipset.go

@@ -0,0 +1,199 @@
+// Copyright 2019 the Kilo authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ipset
+
+import (
+	"bytes"
+	"fmt"
+	"net"
+	"os/exec"
+	"sync"
+	"time"
+)
+
+// Set represents an ipset.
+// Set can safely be used concurrently.
+type Set struct {
+	errors     chan error
+	hosts      map[string]struct{}
+	mu         sync.Mutex
+	name       string
+	subscribed bool
+
+	// Make these functions fields to allow
+	// for testing.
+	add func(string) error
+	del func(string) error
+}
+
+func setExists(name string) (bool, error) {
+	cmd := exec.Command("ipset", "list", "-n")
+	var stderr, stdout bytes.Buffer
+	cmd.Stderr = &stderr
+	cmd.Stdout = &stdout
+	if err := cmd.Run(); err != nil {
+		return false, fmt.Errorf("failed to check for set %s: %s", name, stderr.String())
+	}
+	return bytes.Contains(stdout.Bytes(), []byte(name)), nil
+}
+
+func hostInSet(set, name string) (bool, error) {
+	cmd := exec.Command("ipset", "list", set)
+	var stderr, stdout bytes.Buffer
+	cmd.Stderr = &stderr
+	cmd.Stdout = &stdout
+	if err := cmd.Run(); err != nil {
+		return false, fmt.Errorf("failed to check for host %s: %s", name, stderr.String())
+	}
+	return bytes.Contains(stdout.Bytes(), []byte(name)), nil
+}
+
+// New generates a new ipset.
+func New(name string) *Set {
+	return &Set{
+		errors: make(chan error),
+		hosts:  make(map[string]struct{}),
+		name:   name,
+
+		add: func(ip string) error {
+			ok, err := hostInSet(name, ip)
+			if err != nil {
+				return err
+			}
+			if !ok {
+				cmd := exec.Command("ipset", "add", name, ip)
+				var stderr bytes.Buffer
+				cmd.Stderr = &stderr
+				if err := cmd.Run(); err != nil {
+					return fmt.Errorf("failed to add host %s to set %s: %s", ip, name, stderr.String())
+				}
+			}
+			return nil
+		},
+		del: func(ip string) error {
+			ok, err := hostInSet(name, ip)
+			if err != nil {
+				return err
+			}
+			if ok {
+				cmd := exec.Command("ipset", "del", name, ip)
+				var stderr bytes.Buffer
+				cmd.Stderr = &stderr
+				if err := cmd.Run(); err != nil {
+					return fmt.Errorf("failed to remove host %s from set %s: %s", ip, name, stderr.String())
+				}
+			}
+			return nil
+		},
+	}
+}
+
+// Run watches for changes to the ipset and reconciles
+// the ipset against the desired state.
+func (s *Set) Run(stop <-chan struct{}) (<-chan error, error) {
+	s.mu.Lock()
+	if s.subscribed {
+		s.mu.Unlock()
+		return s.errors, nil
+	}
+	// Ensure a given instance only subscribes once.
+	s.subscribed = true
+	s.mu.Unlock()
+	go func() {
+		defer close(s.errors)
+		for {
+			select {
+			case <-time.After(2 * time.Second):
+			case <-stop:
+				return
+			}
+			ok, err := setExists(s.name)
+			if err != nil {
+				nonBlockingSend(s.errors, err)
+			}
+			// The set does not exist so wait and try again later.
+			if !ok {
+				continue
+			}
+			s.mu.Lock()
+			for h := range s.hosts {
+				if err := s.add(h); err != nil {
+					nonBlockingSend(s.errors, err)
+				}
+			}
+			s.mu.Unlock()
+		}
+	}()
+	return s.errors, nil
+}
+
+// CleanUp will clean up any hosts added to the set.
+func (s *Set) CleanUp() error {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	for h := range s.hosts {
+		if err := s.del(h); err != nil {
+			return err
+		}
+		delete(s.hosts, h)
+	}
+	return nil
+}
+
+// Set idempotently overwrites any hosts previously defined
+// for the ipset with the given hosts.
+func (s *Set) Set(hosts []net.IP) error {
+	h := make(map[string]struct{})
+	for _, host := range hosts {
+		if host == nil {
+			continue
+		}
+		h[host.String()] = struct{}{}
+	}
+	exists, err := setExists(s.name)
+	if err != nil {
+		return err
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	for k := range s.hosts {
+		if _, ok := h[k]; !ok {
+			if exists {
+				if err := s.del(k); err != nil {
+					return err
+				}
+			}
+			delete(s.hosts, k)
+		}
+	}
+	for k := range h {
+		if _, ok := s.hosts[k]; !ok {
+			if exists {
+				if err := s.add(k); err != nil {
+					return err
+				}
+			}
+			s.hosts[k] = struct{}{}
+		}
+	}
+	return nil
+}
+
+func nonBlockingSend(errors chan<- error, err error) {
+	select {
+	case errors <- err:
+	default:
+	}
+}