github.com/keybase/client/go@v0.0.0-20241007131713-f10651d043c8/chat/attachments/stash.go (about)

     1  package attachments
     2  
     3  import (
     4  	"encoding/gob"
     5  	"errors"
     6  	"fmt"
     7  	"os"
     8  	"path/filepath"
     9  	"sync"
    10  	"time"
    11  
    12  	"github.com/keybase/client/go/chat/signencrypt"
    13  	"github.com/keybase/client/go/protocol/chat1"
    14  )
    15  
    16  type AttachmentInfo struct {
    17  	ObjectKey string                   // s3 destination
    18  	EncKey    signencrypt.SecretboxKey // encryption key
    19  	SignKey   signencrypt.SignKey      // signing key
    20  	VerifyKey signencrypt.VerifyKey    // verification key
    21  	Parts     map[int]string           // map of parts uploaded to S3, key == part number, value == hash of ciphertext
    22  	StartedAt time.Time                // when the upload started
    23  }
    24  
    25  type StashKey struct {
    26  	OutboxID chat1.OutboxID
    27  	Preview  bool
    28  }
    29  
    30  func NewStashKey(outboxID chat1.OutboxID, preview bool) StashKey {
    31  	return StashKey{
    32  		OutboxID: outboxID,
    33  		Preview:  preview,
    34  	}
    35  }
    36  
    37  func (s StashKey) String() string {
    38  	return fmt.Sprintf("%s:%v", s.OutboxID, s.Preview)
    39  }
    40  
    41  type AttachmentStash interface {
    42  	Start(key StashKey, info AttachmentInfo) error
    43  	Lookup(key StashKey) (AttachmentInfo, bool, error)
    44  	RecordPart(key StashKey, partNumber int, hash string) error
    45  	Finish(key StashKey) error
    46  }
    47  
    48  var ErrPartNotFound = errors.New("part does not exist in stash")
    49  
    50  type FileStash struct {
    51  	dir string
    52  	sync.Mutex
    53  }
    54  
    55  func NewFileStash(dir string) *FileStash {
    56  	return &FileStash{dir: dir}
    57  }
    58  
    59  func (f *FileStash) Start(key StashKey, info AttachmentInfo) error {
    60  	f.Lock()
    61  	defer f.Unlock()
    62  	c, err := f.contents()
    63  	if err != nil {
    64  		return err
    65  	}
    66  	info.StartedAt = time.Now()
    67  	c[key.String()] = info
    68  
    69  	return f.serialize(c)
    70  }
    71  
    72  func (f *FileStash) Lookup(key StashKey) (AttachmentInfo, bool, error) {
    73  	f.Lock()
    74  	defer f.Unlock()
    75  	return f.lookup(key)
    76  }
    77  
    78  func (f *FileStash) RecordPart(key StashKey, partNumber int, hash string) error {
    79  	f.Lock()
    80  	defer f.Unlock()
    81  	c, err := f.contents()
    82  	if err != nil {
    83  		return err
    84  	}
    85  	info, found := c[key.String()]
    86  	if !found {
    87  		return ErrPartNotFound
    88  	}
    89  
    90  	if info.Parts == nil {
    91  		info.Parts = make(map[int]string)
    92  	}
    93  
    94  	info.Parts[partNumber] = hash
    95  	c[key.String()] = info
    96  	return f.serialize(c)
    97  }
    98  
    99  func (f *FileStash) Finish(key StashKey) error {
   100  	f.Lock()
   101  	defer f.Unlock()
   102  	c, err := f.contents()
   103  	if err != nil {
   104  		return err
   105  	}
   106  	delete(c, key.String())
   107  	return f.serialize(c)
   108  }
   109  
   110  func (f *FileStash) filename() string {
   111  	if f.dir == "" {
   112  		panic("FileStash used with no directory")
   113  	}
   114  	return filepath.Join(f.dir, "chat_attachment_stash")
   115  }
   116  
   117  func (f *FileStash) contents() (map[string]AttachmentInfo, error) {
   118  	x, err := os.Open(f.filename())
   119  	if err != nil {
   120  		if os.IsNotExist(err) {
   121  			return make(map[string]AttachmentInfo), nil
   122  		}
   123  		return nil, err
   124  	}
   125  	defer x.Close()
   126  
   127  	v := make(map[string]AttachmentInfo)
   128  	dec := gob.NewDecoder(x)
   129  	if err := dec.Decode(&v); err != nil {
   130  		return nil, err
   131  	}
   132  	return v, nil
   133  }
   134  
   135  func (f *FileStash) serialize(m map[string]AttachmentInfo) error {
   136  	x, err := os.Create(f.filename())
   137  	if err != nil {
   138  		return err
   139  	}
   140  	defer x.Close()
   141  	enc := gob.NewEncoder(x)
   142  	return enc.Encode(m)
   143  }
   144  
   145  func (f *FileStash) lookup(key StashKey) (AttachmentInfo, bool, error) {
   146  	c, err := f.contents()
   147  	if err != nil {
   148  		return AttachmentInfo{}, false, err
   149  	}
   150  	info, found := c[key.String()]
   151  	return info, found, nil
   152  }