Skip to content

Commit

Permalink
add reusable buffer for transactions
Browse files Browse the repository at this point in the history
Signed-off-by: Daman Arora <aroradaman@gmail.com>
  • Loading branch information
aroradaman authored and Daman Arora committed May 20, 2024
1 parent 5951417 commit 841e0a5
Show file tree
Hide file tree
Showing 2 changed files with 23 additions and 12 deletions.
25 changes: 19 additions & 6 deletions nftables.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,11 +17,13 @@ limitations under the License.
package knftables

import (
"bytes"
"context"
"encoding/json"
"fmt"
"os/exec"
"strings"
"sync"
)

// Interface is an interface for running nftables commands against a given family and table.
Expand Down Expand Up @@ -71,6 +73,9 @@ type nftContext struct {
type realNFTables struct {
nftContext

bufferMutex sync.Mutex
buffer *bytes.Buffer

exec execer
path string
}
Expand All @@ -85,8 +90,8 @@ func newInternal(family Family, table string, execer execer) (Interface, error)
family: family,
table: table,
},

exec: execer,
buffer: &bytes.Buffer{},
exec: execer,
}

nft.path, err = nft.exec.LookPath("nft")
Expand Down Expand Up @@ -136,34 +141,42 @@ func (nft *realNFTables) NewTransaction() *Transaction {

// Run is part of Interface
func (nft *realNFTables) Run(ctx context.Context, tx *Transaction) error {
nft.bufferMutex.Lock()
defer nft.bufferMutex.Unlock()

if tx.err != nil {
return tx.err
}

buf, err := tx.asCommandBuf()
nft.buffer.Reset()
err := tx.populateCommandBuf(nft.buffer)
if err != nil {
return err
}

cmd := exec.CommandContext(ctx, nft.path, "-f", "-")
cmd.Stdin = buf
cmd.Stdin = nft.buffer
_, err = nft.exec.Run(cmd)
return err
}

// Check is part of Interface
func (nft *realNFTables) Check(ctx context.Context, tx *Transaction) error {
nft.bufferMutex.Lock()
defer nft.bufferMutex.Unlock()

if tx.err != nil {
return tx.err
}

buf, err := tx.asCommandBuf()
nft.buffer.Reset()
err := tx.populateCommandBuf(nft.buffer)
if err != nil {
return err
}

cmd := exec.CommandContext(ctx, nft.path, "--check", "-f", "-")
cmd.Stdin = buf
cmd.Stdin = nft.buffer
_, err = nft.exec.Run(cmd)
return err
}
Expand Down
10 changes: 4 additions & 6 deletions transaction.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,7 +19,6 @@ package knftables
import (
"bytes"
"fmt"
"io"
)

// Transaction represents an nftables transaction
Expand Down Expand Up @@ -48,17 +47,16 @@ const (
flushVerb verb = "flush"
)

// asCommandBuf returns the transaction as an io.Reader that outputs a series of nft commands
func (tx *Transaction) asCommandBuf() (io.Reader, error) {
// populateCommandBuf populates the transaction as series of nft commands to the given bytes.Buffer.
func (tx *Transaction) populateCommandBuf(buf *bytes.Buffer) error {
if tx.err != nil {
return nil, tx.err
return tx.err
}

buf := &bytes.Buffer{}
for _, op := range tx.operations {
op.obj.writeOperation(op.verb, tx.nftContext, buf)
}
return buf, nil
return nil
}

// String returns the transaction as a string containing the nft commands; if there is
Expand Down

0 comments on commit 841e0a5

Please sign in to comment.