forked from projectcalico/felix
-
Notifications
You must be signed in to change notification settings - Fork 3
/
restore_buffer.go
134 lines (120 loc) · 4.55 KB
/
restore_buffer.go
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
123
124
125
126
127
128
129
130
131
132
133
134
// Copyright (c) 2019 Tigera, Inc. All rights reserved.
//
// 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 iptables
import (
"bytes"
"fmt"
log "github.com/sirupsen/logrus"
)
// RestoreInputBuilder builds a byte slice for use as input to iptables-restore.
//
// Operations must be done inside a per-table transaction.
//
// Example:
//
// buf.Reset() // Reset the buffer, if needed.
// buf.StartTransaction("filter")
// buf.WriteForwardReference("cali-chain-name")
// buf.WriteLine("-A cali-chain-name ...")
// buf.EndTransaction()
// bytes = buf.GetBytesAndReset()
// <write bytes to iptables-restore stdin>
//
// Transactions are ignored completely if there are no writes between the StartTransaction()
// and EndTransaction() calls.
type RestoreInputBuilder struct {
buf bytes.Buffer
currentTableName string
txnOpenerWritten bool
NumLinesWritten counter
}
// Empty returns true if there is nothing in the buffer (i.e. all the transactions stored in the buffer were no-ops).
func (b *RestoreInputBuilder) Empty() bool {
return b.buf.Len() == 0
}
// Reset the builder completely, any pending transaction is discarded.
func (b *RestoreInputBuilder) Reset() {
b.buf.Reset()
b.currentTableName = ""
b.txnOpenerWritten = false
}
// StartTransaction opens a new transaction context for the named table.
// Panics if there is already a transaction in progress.
func (b *RestoreInputBuilder) StartTransaction(tableName string) {
if b.currentTableName != "" {
log.Panic("StartTransaction() called without ending previous transaction.")
}
b.currentTableName = tableName
b.txnOpenerWritten = false
}
// EndTransaction ends the open transaction, if the transaction was non-empty, writes a COMMIT.
// Resets the transaction tracking state. Panics if there was no open transaction.
func (b *RestoreInputBuilder) EndTransaction() {
if b.currentTableName == "" {
log.Panic("EndTransaction() called without active transaction.")
}
if b.txnOpenerWritten {
b.writeFormattedLine("COMMIT")
}
b.currentTableName = ""
}
// writeFormattedLine writes a line to the internal buffer, appending a new line.
func (b *RestoreInputBuilder) writeFormattedLine(format string, args ...interface{}) {
_, err := fmt.Fprintf(&b.buf, format, args...)
if err != nil {
log.WithError(err).Panic("Failed to write to in-memory buffer")
}
b.buf.WriteString("\n")
if b.NumLinesWritten != nil {
b.NumLinesWritten.Inc()
}
}
// maybeWriteTransactionOpener ensures that the transaction opening line has been written.
// Panics if there is no open transaction.
func (b *RestoreInputBuilder) maybeWriteTransactionOpener() {
if b.currentTableName == "" {
log.Panic("maybeWriteTransactionOpener() called without active transaction.")
}
if !b.txnOpenerWritten {
b.writeFormattedLine("*%s", b.currentTableName)
b.txnOpenerWritten = true
}
}
// WriteForwardReference writes a "forward reference" for the given chain name. A forward reference is an instruction
// that tells iptables to ensure that the given chain exists and that it is empty. Panics if there is no open
// transaction.
func (b *RestoreInputBuilder) WriteForwardReference(chainName string) {
b.maybeWriteTransactionOpener()
b.writeFormattedLine(":%s - -", chainName)
}
// WriteLine writes a line of iptables instructions to the buffer. Intended for writing the actual rules.
// Panics if there is no open transaction.
func (b *RestoreInputBuilder) WriteLine(line string) {
b.maybeWriteTransactionOpener()
b.writeFormattedLine(line)
}
// GetBytesAndReset returns the contents of the buffer and, as a side effect, resets the buffer. For performance,
// this is a direct reference to the data rather than a copy. The returned slice is only valid until the next
// write operation on the builder. Should be called after EndTransaction; panics if there is a still-open transaction.
func (b *RestoreInputBuilder) GetBytesAndReset() []byte {
if b.currentTableName != "" {
log.Panic("GetBytesAndReset() called inside transaction.")
}
buf := b.buf.Next(b.buf.Len())
b.Reset()
return buf
}
type counter interface {
Inc()
}