1 files changed, 25 insertions, 33 deletions

diff --git a/roseh.moe.go b/roseh.moe.go
index e6db419..f9c5478 100644
--- a/roseh.moe.go
+++ b/roseh.moe.go
@@ -254,15 +254,16 @@ func login(w http.ResponseWriter, r *http.Request) {
 }
 
 func readNotepad(key cryptoutil.EncryptionKey) (string, error) {
-	encrypted, err := os.ReadFile(*notepadDir + "/notepad")
+	f, err := os.Open(*notepadDir + "/notepad")
 	if err != nil {
 		return "", err
 	}
-	decrypted, err := key.Decrypt(encrypted)
+	defer f.Close()
+	notepad, err := io.ReadAll(key.NewReader(f, []byte("notepad")))
 	if err != nil {
 		return "", err
 	}
-	return string(decrypted), nil
+	return string(notepad), nil
 }
 
 var (
@@ -315,16 +316,17 @@ func saveNote(w http.ResponseWriter, r *http.Request) error {
 	if key == nil {
 		return errors.New("not logged in")
 	}
-	encrypted, err := key.Encrypt([]byte(r.FormValue("content")))
-	if err != nil {
-		return err
-	}
 	f, err := os.CreateTemp(*notepadDir, "notepad")
 	if err != nil {
 		return err
 	}
 	defer f.Close()
-	if _, err = f.Write(encrypted); err != nil {
+	encryptingWriter := key.NewWriter(f, []byte("notepad"))
+	if _, err := encryptingWriter.Write([]byte(r.FormValue("content"))); err != nil {
+		os.Remove(f.Name())
+		return err
+	}
+	if err := encryptingWriter.Close(); err != nil {
 		os.Remove(f.Name())
 		return err
 	}
@@ -434,10 +436,7 @@ func createNote(stream *gob.Decoder) (*api.CreateNoteResponse, error) {
 		break
 	}
 	defer f.Close()
-	encryptingWriter, err := key.Writer(f)
-	if err != nil {
-		return nil, err
-	}
+	encryptingWriter := key.NewWriter(f, []byte("notes/"+name))
 	for {
 		req := new(api.CreateNoteRequestStream)
 		if err := stream.Decode(req); err != nil {
@@ -450,7 +449,7 @@ func createNote(stream *gob.Decoder) (*api.CreateNoteResponse, error) {
 			return nil, fmt.Errorf("create note: write note to file: %s", err)
 		}
 	}
-	if err := encryptingWriter.Flush(); err != nil {
+	if err := encryptingWriter.Close(); err != nil {
 		return nil, err
 	}
 	if err := f.Close(); err != nil {
@@ -459,6 +458,17 @@ func createNote(stream *gob.Decoder) (*api.CreateNoteResponse, error) {
 	return &api.CreateNoteResponse{Name: name}, nil
 }
 
+type readNoteResponseWriter struct {
+	stream *encoder
+}
+
+func (w *readNoteResponseWriter) Write(buf []byte) (int, error) {
+	if err := w.stream.send(&api.ReadNoteResponseStream{Chunk: buf}); err != nil {
+		return 0, err
+	}
+	return len(buf), nil
+}
+
 func readNote(stream *encoder, req *api.ReadNoteRequest) error {
 	f, err := os.OpenInRoot(*notepadDir+"/notes", req.Note)
 	if err != nil {
@@ -474,26 +484,8 @@ func readNote(stream *encoder, req *api.ReadNoteRequest) error {
 	if key == nil {
 		return fmt.Errorf("%w: need login", api.PermissionDenied)
 	}
-	decryptingReader, err := key.Reader(f)
-	if err != nil {
-		return err
-	}
-	buf := make([]byte, 4*1024*1024)
-	for {
-		n, err := decryptingReader.Read(buf)
-		if n > 0 {
-			if err := stream.send(&api.ReadNoteResponseStream{Chunk: buf[:n]}); err != nil {
-				return err
-			}
-		}
-		if err != nil {
-			if errors.Is(err, io.EOF) {
-				break
-			}
-			return err
-		}
-	}
-	return nil
+	_, err = io.Copy(&readNoteResponseWriter{stream: stream}, key.NewReader(f, []byte("notes/"+req.Note)))
+	return err
 }
 
 func tokenAuth(w http.ResponseWriter, r *http.Request) bool {