aboutsummaryrefslogtreecommitdiff
path: root/pkg/atlas/util.go
blob: 7ea8b1439bf5a8ee89b23a8d0fbedcaad68428d8 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
package atlas

import (
	"fmt"
	"io"
	"net/http"
	"net/netip"
	"os"
	"sync"

	"github.com/pg9182/ip2x"
	"github.com/rs/zerolog"
)

// ip2xMgr wraps a file-backed IP2Location database.
type ip2xMgr struct {
	file *os.File
	db   *ip2x.DB
	mu   sync.RWMutex
}

// Load replaces the currently loaded database with the specified file. If name
// is empty, the existing database, if any, is reopened.
func (m *ip2xMgr) Load(name string) error {
	if name == "" {
		m.mu.RLock()
		if m.file == nil {
			return fmt.Errorf("no ip2location database loaded")
		}
		name = m.file.Name()
		m.mu.RUnlock()
	}

	f, err := os.Open(name)
	if err != nil {
		return err
	}

	db, err := ip2x.New(f)
	if err != nil {
		f.Close()
		return err
	}

	if p, _ := db.Info(); p != ip2x.IP2Location {
		f.Close()
		return fmt.Errorf("not an ip2location database")
	}

	m.mu.Lock()
	defer m.mu.Unlock()

	m.file.Close()
	m.file = f
	m.db = db
	return nil
}

// Lookup calls [ip2x.DB.Lookup] if a database is loaded.
func (m *ip2xMgr) LookupFields(ip netip.Addr) (ip2x.Record, error) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	if m.db == nil {
		return ip2x.Record{}, fmt.Errorf("no ip2location database loaded")
	}
	return m.db.Lookup(ip)
}

type zerologWriterLevel struct {
	w io.Writer // or zerolog.LevelWriter
	l zerolog.Level
	m sync.Mutex
}

var _ zerolog.LevelWriter = (*zerologWriterLevel)(nil)

func newZerologWriterLevel(w io.Writer, l zerolog.Level) *zerologWriterLevel {
	return &zerologWriterLevel{w: w, l: l}
}

func (wl *zerologWriterLevel) Write(p []byte) (n int, err error) {
	wl.m.Lock()
	defer wl.m.Unlock()
	if wl.w != nil {
		return wl.w.Write(p)
	}
	return len(p), nil
}

func (wl *zerologWriterLevel) WriteLevel(l zerolog.Level, p []byte) (n int, err error) {
	if l >= wl.l {
		wl.m.Lock()
		defer wl.m.Unlock()
		if wl.w != nil {
			if lw, ok := wl.w.(zerolog.LevelWriter); ok {
				return lw.WriteLevel(l, p)
			}
			return wl.w.Write(p)
		}
	}
	return len(p), nil
}

func (wl *zerologWriterLevel) SwapWriter(fn func(io.Writer) io.Writer) {
	wl.m.Lock()
	defer wl.m.Unlock()
	wl.w = fn(wl.w)
}

type middlewares []func(http.Handler) http.Handler

func (ms *middlewares) Add(m func(http.Handler) http.Handler) *middlewares {
	*ms = append(*ms, m)
	return ms
}

func (ms *middlewares) Then(h http.Handler) http.Handler {
	for i := len(*ms) - 1; i >= 0; i-- {
		h = (*ms)[i](h)
	}
	return h
}