219 lines
6.9 KiB
Go
219 lines
6.9 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"database/sql"
|
|
"errors"
|
|
"linkport/internal/model"
|
|
"linkport/internal/security"
|
|
)
|
|
|
|
const commandOutputChunkBytes = 32 << 10
|
|
|
|
func hasCommandResult(protocol, kind string) bool {
|
|
return protocol == "ssh" && (kind == "shell_input" || kind == "exec") || protocol == "mysql" && (kind == "query" || kind == "prepared")
|
|
}
|
|
|
|
func (s *Store) RecordCommandOutput(p model.Proxy, source string, output model.CommandOutput) error {
|
|
validStream := p.Type == "ssh" && (output.Stream == "stdout" || output.Stream == "stderr") || p.Type == "mysql" && output.Stream == "result"
|
|
if (p.Type != "ssh" && p.Type != "mysql") || output.CommandID == "" || output.RecordedAt.IsZero() || (len(output.Data) > 0 && !validStream) {
|
|
return errors.New("invalid command output audit")
|
|
}
|
|
s.auditOps.Lock()
|
|
defer s.auditOps.Unlock()
|
|
// Keep a command and all of its result/output rows in one shard. If a
|
|
// single command exceeds the configured shard limit, the next command
|
|
// rotates the file instead of splitting this atomic audit record.
|
|
db, err := s.commandOutputDB(context.Background(), output.CommandID)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
tx, err := db.Begin()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
defer tx.Rollback()
|
|
var id, kind string
|
|
err = tx.QueryRow("SELECT id,kind FROM commands WHERE id=? AND proxy_id=? AND source=? AND protocol=?", output.CommandID, p.ID, source, p.Type).Scan(&id, &kind)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !hasCommandResult(p.Type, kind) {
|
|
return errors.New("unsupported command output kind")
|
|
}
|
|
resultBlob, err := s.seal(output.Result, "command-result:"+id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if _, err = tx.Exec("INSERT INTO command_results(command_id,content) VALUES(?,?) ON CONFLICT(command_id) DO UPDATE SET content=excluded.content", id, resultBlob); err != nil {
|
|
return err
|
|
}
|
|
for data := output.Data; len(data) > 0; {
|
|
size := min(len(data), commandOutputChunkBytes)
|
|
chunkID := security.Token()
|
|
blob, err := s.vault.Seal(data[:size], outputPurpose(id, chunkID, output.Stream))
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if _, err = tx.Exec("INSERT INTO command_outputs(id,command_id,stream,content,size) VALUES(?,?,?,?,?)", chunkID, id, output.Stream, blob, size); err != nil {
|
|
return err
|
|
}
|
|
data = data[size:]
|
|
}
|
|
return tx.Commit()
|
|
}
|
|
|
|
func outputPurpose(command, chunk, stream string) string {
|
|
return "command-output:" + command + ":" + chunk + ":" + stream
|
|
}
|
|
|
|
type CommandOutputChunk struct {
|
|
Seq int64 `json:"seq"`
|
|
Stream string `json:"stream"`
|
|
Data []byte `json:"data"` // Base64 preserves arbitrary bytes and split UTF-8.
|
|
}
|
|
|
|
type CommandOutputPage struct {
|
|
Result *model.CommandResult `json:"result"`
|
|
Items []CommandOutputChunk `json:"items"`
|
|
Bytes int64 `json:"bytes"`
|
|
NextAfter int64 `json:"nextAfter"`
|
|
HasMore bool `json:"hasMore"`
|
|
}
|
|
|
|
// Caller holds auditOps. Never nest this lookup inside a live command-row walk:
|
|
// each shard intentionally has only one database connection.
|
|
func (s *Store) commandOutputDB(ctx context.Context, id string) (*sql.DB, error) {
|
|
if s.closed {
|
|
return nil, errors.New("store is closed")
|
|
}
|
|
days, err := s.auditDays()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
for _, day := range days {
|
|
db, err := s.auditDB(day)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
var found string
|
|
err = db.QueryRowContext(ctx, "SELECT id FROM commands WHERE id=?", id).Scan(&found)
|
|
if err == nil {
|
|
return db, nil
|
|
}
|
|
if !errors.Is(err, sql.ErrNoRows) {
|
|
return nil, err
|
|
}
|
|
}
|
|
return nil, sql.ErrNoRows
|
|
}
|
|
|
|
func (s *Store) CommandOutput(ctx context.Context, id, stream string, after int64) (CommandOutputPage, error) {
|
|
s.auditOps.RLock()
|
|
defer s.auditOps.RUnlock()
|
|
page := CommandOutputPage{Items: []CommandOutputChunk{}, NextAfter: after}
|
|
db, err := s.commandOutputDB(ctx, id)
|
|
if err != nil {
|
|
return page, err
|
|
}
|
|
// Keep metadata, size and chunks in one consistent snapshot.
|
|
tx, err := db.BeginTx(ctx, &sql.TxOptions{ReadOnly: true})
|
|
if err != nil {
|
|
return page, err
|
|
}
|
|
defer tx.Rollback()
|
|
var blob []byte
|
|
err = tx.QueryRowContext(ctx, "SELECT content FROM command_results WHERE command_id=?", id).Scan(&blob)
|
|
if errors.Is(err, sql.ErrNoRows) {
|
|
return page, nil
|
|
}
|
|
if err != nil {
|
|
return page, err
|
|
}
|
|
page.Result = &model.CommandResult{}
|
|
if err = s.unseal(blob, "command-result:"+id, page.Result); err != nil {
|
|
return page, err
|
|
}
|
|
if err = tx.QueryRowContext(ctx, "SELECT COALESCE(SUM(size),0) FROM command_outputs WHERE command_id=? AND (?='' OR stream=?)", id, stream, stream).Scan(&page.Bytes); err != nil {
|
|
return page, err
|
|
}
|
|
rows, err := tx.QueryContext(ctx, "SELECT seq,id,stream,content FROM command_outputs WHERE command_id=? AND seq>? AND (?='' OR stream=?) ORDER BY seq LIMIT 33", id, after, stream, stream)
|
|
if err != nil {
|
|
return page, err
|
|
}
|
|
defer rows.Close()
|
|
for rows.Next() {
|
|
if len(page.Items) == 32 {
|
|
page.HasMore = true
|
|
break
|
|
}
|
|
var chunk CommandOutputChunk
|
|
var chunkID string
|
|
if err = rows.Scan(&chunk.Seq, &chunkID, &chunk.Stream, &blob); err != nil {
|
|
return page, err
|
|
}
|
|
chunk.Data, err = s.vault.Open(blob, outputPurpose(id, chunkID, chunk.Stream))
|
|
if err != nil {
|
|
return page, err
|
|
}
|
|
page.Items = append(page.Items, chunk)
|
|
page.NextAfter = chunk.Seq
|
|
}
|
|
return page, rows.Err()
|
|
}
|
|
|
|
// Download is bounded to the chunks present at the start. Holding auditOps
|
|
// prevents retention from deleting its shard while the response is streaming.
|
|
func (s *Store) DownloadCommandOutput(ctx context.Context, id, stream string, begin func(int64), write func([]byte) error) error {
|
|
s.auditOps.RLock()
|
|
defer s.auditOps.RUnlock()
|
|
db, err := s.commandOutputDB(ctx, id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
var maxSeq, size int64
|
|
err = db.QueryRowContext(ctx, "SELECT COALESCE(MAX(seq),0),COALESCE(SUM(size),0) FROM command_outputs WHERE command_id=? AND (?='' OR stream=?)", id, stream, stream).Scan(&maxSeq, &size)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
begin(size)
|
|
// Release rows/connection between pages so downloads do not starve capture.
|
|
for after := int64(0); after < maxSeq; {
|
|
rows, err := db.QueryContext(ctx, "SELECT seq,id,stream,content FROM command_outputs WHERE command_id=? AND seq>? AND seq<=? AND (?='' OR stream=?) ORDER BY seq LIMIT 32", id, after, maxSeq, stream, stream)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
chunks := make([][]byte, 0, 32)
|
|
for rows.Next() {
|
|
var chunkID, channel string
|
|
var blob []byte
|
|
if err = rows.Scan(&after, &chunkID, &channel, &blob); err != nil {
|
|
break
|
|
}
|
|
var data []byte
|
|
data, err = s.vault.Open(blob, outputPurpose(id, chunkID, channel))
|
|
if err != nil {
|
|
break
|
|
}
|
|
chunks = append(chunks, data)
|
|
}
|
|
rowErr := rows.Err()
|
|
rows.Close()
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if rowErr != nil {
|
|
return rowErr
|
|
}
|
|
if len(chunks) == 0 {
|
|
return errors.New("output snapshot is incomplete")
|
|
}
|
|
for _, data := range chunks {
|
|
if err = write(data); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
}
|
|
return nil
|
|
}
|