Refactor: pngmeta rewrite

This commit is contained in:
Grail Finder
2025-03-10 20:08:07 +03:00
parent 10f0efbb2a
commit 0497cdcfe6
5 changed files with 410 additions and 111 deletions

137
pngmeta/altwriter.go Normal file
View File

@@ -0,0 +1,137 @@
package pngmeta
import (
"bytes"
"elefant/models"
"encoding/base64"
"encoding/binary"
"encoding/json"
"errors"
"fmt"
"hash/crc32"
"io"
"os"
)
const (
pngHeader = "\x89PNG\r\n\x1a\n"
textChunkType = "tEXt"
)
// // PngEmbed holds the key-value pair to embed into the PNG file.
// type PngEmbed struct {
// Key string
// Value string
// }
// WriteToPng embeds the metadata into the specified PNG file and writes the result to outfile.
func WriteToPng(metadata *models.CharCardSpec, sourcePath, outfile string) error {
pngData, err := os.ReadFile(sourcePath)
if err != nil {
return err
}
jsonData, err := json.Marshal(metadata)
if err != nil {
return err
}
base64Data := base64.StdEncoding.EncodeToString(jsonData)
embedData := PngEmbed{
Key: "elefant", // Replace with appropriate key constant
Value: base64Data,
}
var outputBuffer bytes.Buffer
if _, err := outputBuffer.Write([]byte(pngHeader)); err != nil {
return err
}
chunks, iend, err := processChunks(pngData[8:])
if err != nil {
return err
}
for _, chunk := range chunks {
outputBuffer.Write(chunk)
}
newChunk := createTextChunk(embedData)
outputBuffer.Write(newChunk)
outputBuffer.Write(iend)
return os.WriteFile(outfile, outputBuffer.Bytes(), 0666)
}
// processChunks extracts non-tEXt chunks and locates the IEND chunk
func processChunks(data []byte) ([][]byte, []byte, error) {
var (
chunks [][]byte
iendChunk []byte
reader = bytes.NewReader(data)
)
for {
var chunkLength uint32
if err := binary.Read(reader, binary.BigEndian, &chunkLength); err != nil {
if errors.Is(err, io.EOF) {
break
}
return nil, nil, fmt.Errorf("error reading chunk length: %w", err)
}
chunkType := make([]byte, 4)
if _, err := reader.Read(chunkType); err != nil {
return nil, nil, fmt.Errorf("error reading chunk type: %w", err)
}
chunkData := make([]byte, chunkLength)
if _, err := reader.Read(chunkData); err != nil {
return nil, nil, fmt.Errorf("error reading chunk data: %w", err)
}
crc := make([]byte, 4)
if _, err := reader.Read(crc); err != nil {
return nil, nil, fmt.Errorf("error reading CRC: %w", err)
}
fullChunk := bytes.NewBuffer(nil)
binary.Write(fullChunk, binary.BigEndian, chunkLength)
fullChunk.Write(chunkType)
fullChunk.Write(chunkData)
fullChunk.Write(crc)
switch string(chunkType) {
case "IEND":
iendChunk = fullChunk.Bytes()
return chunks, iendChunk, nil
case textChunkType:
continue // Skip existing tEXt chunks
default:
chunks = append(chunks, fullChunk.Bytes())
}
}
return nil, nil, errors.New("IEND chunk not found")
}
// createTextChunk generates a valid tEXt chunk with proper CRC
func createTextChunk(embed PngEmbed) []byte {
content := bytes.NewBuffer(nil)
content.WriteString(embed.Key)
content.WriteByte(0) // Null separator
content.WriteString(embed.Value)
data := content.Bytes()
crc := crc32.NewIEEE()
crc.Write([]byte(textChunkType))
crc.Write(data)
chunk := bytes.NewBuffer(nil)
binary.Write(chunk, binary.BigEndian, uint32(len(data)))
chunk.Write([]byte(textChunkType))
chunk.Write(data)
binary.Write(chunk, binary.BigEndian, crc.Sum32())
return chunk.Bytes()
}

View File

@@ -1,7 +1,19 @@
package pngmeta
import (
"bytes"
"elefant/models"
"encoding/base64"
"encoding/binary"
"encoding/json"
"errors"
"fmt"
"image"
"image/color"
"image/png"
"io"
"os"
"path/filepath"
"testing"
)
@@ -28,3 +40,155 @@ func TestReadMeta(t *testing.T) {
})
}
}
// Test helper: Create a simple PNG image with test shapes
func createTestImage(t *testing.T) string {
img := image.NewRGBA(image.Rect(0, 0, 200, 200))
// Fill background with white
for y := 0; y < 200; y++ {
for x := 0; x < 200; x++ {
img.Set(x, y, color.White)
}
}
// Draw a red square
for y := 50; y < 150; y++ {
for x := 50; x < 150; x++ {
img.Set(x, y, color.RGBA{R: 255, A: 255})
}
}
// Draw a blue circle
center := image.Point{100, 100}
radius := 40
for y := center.Y - radius; y <= center.Y+radius; y++ {
for x := center.X - radius; x <= center.X+radius; x++ {
dx := x - center.X
dy := y - center.Y
if dx*dx+dy*dy <= radius*radius {
img.Set(x, y, color.RGBA{B: 255, A: 255})
}
}
}
// Create temp file
tmpDir := t.TempDir()
fpath := filepath.Join(tmpDir, "test-image.png")
f, err := os.Create(fpath)
if err != nil {
t.Fatalf("Error creating temp file: %v", err)
}
defer f.Close()
if err := png.Encode(f, img); err != nil {
t.Fatalf("Error encoding PNG: %v", err)
}
return fpath
}
func TestWriteToPng(t *testing.T) {
// Create test image
srcPath := createTestImage(t)
dstPath := filepath.Join(filepath.Dir(srcPath), "output.png")
// dstPath := "test.png"
// Create test metadata
metadata := &models.CharCardSpec{
Description: "Test image containing a red square and blue circle on white background",
}
// Embed metadata
if err := WriteToPng(metadata, srcPath, dstPath); err != nil {
t.Fatalf("WriteToPng failed: %v", err)
}
// Verify output file exists
if _, err := os.Stat(dstPath); os.IsNotExist(err) {
t.Fatalf("Output file not created: %v", err)
}
// Read and verify metadata
t.Run("VerifyMetadata", func(t *testing.T) {
data, err := os.ReadFile(dstPath)
if err != nil {
t.Fatalf("Error reading output file: %v", err)
}
// Verify PNG header
if string(data[:8]) != pngHeader {
t.Errorf("Invalid PNG header")
}
// Extract metadata
embedded := extractMetadata(t, data)
if embedded.Description != metadata.Description {
t.Errorf("Metadata mismatch\nWant: %q\nGot: %q",
metadata.Description, embedded.Description)
}
})
// Optional: Add cleanup if needed
// t.Cleanup(func() {
// os.Remove(dstPath)
// })
}
// Helper to extract embedded metadata from PNG bytes
func extractMetadata(t *testing.T, data []byte) *models.CharCardSpec {
r := bytes.NewReader(data[8:]) // Skip PNG header
for {
var length uint32
if err := binary.Read(r, binary.BigEndian, &length); err != nil {
if errors.Is(err, io.EOF) {
break
}
t.Fatalf("Error reading chunk length: %v", err)
}
chunkType := make([]byte, 4)
if _, err := r.Read(chunkType); err != nil {
t.Fatalf("Error reading chunk type: %v", err)
}
// Read chunk data
chunkData := make([]byte, length)
if _, err := r.Read(chunkData); err != nil {
t.Fatalf("Error reading chunk data: %v", err)
}
// Read and discard CRC
if _, err := r.Read(make([]byte, 4)); err != nil {
t.Fatalf("Error reading CRC: %v", err)
}
if string(chunkType) == embType {
parts := bytes.SplitN(chunkData, []byte{0}, 2)
if len(parts) != 2 {
t.Fatalf("Invalid tEXt chunk format")
}
decoded, err := base64.StdEncoding.DecodeString(string(parts[1]))
if err != nil {
t.Fatalf("Base64 decode error: %v", err)
}
var result models.CharCardSpec
if err := json.Unmarshal(decoded, &result); err != nil {
t.Fatalf("JSON unmarshal error: %v", err)
}
return &result
}
}
t.Fatal("Metadata not found in PNG")
return nil
}
func readTextChunk(t *testing.T, r io.ReadSeeker) *models.CharCardSpec {
var length uint32
binary.Read(r, binary.BigEndian, &length)
chunkType := make([]byte, 4)
r.Read(chunkType)
data := make([]byte, length)
r.Read(data)
// Read CRC (but skip validation for test purposes)
crc := make([]byte, 4)
r.Read(crc)
parts := bytes.SplitN(data, []byte{0}, 2) // Split key-value pair
if len(parts) != 2 {
t.Fatalf("Invalid tEXt chunk format")
}
// key := string(parts[0])
value := parts[1]
decoded, err := base64.StdEncoding.DecodeString(string(value))
if err != nil {
t.Fatalf("Base64 decode error: %v; value: %s", err, string(value))
}
var result models.CharCardSpec
if err := json.Unmarshal(decoded, &result); err != nil {
t.Fatalf("JSON unmarshal error: %v", err)
}
return &result
}

View File

@@ -1,116 +1,112 @@
package pngmeta
import (
"bytes"
"elefant/models"
"encoding/base64"
"encoding/binary"
"encoding/json"
"errors"
"fmt"
"hash/crc32"
"io"
"os"
)
// import (
// "bytes"
// "encoding/binary"
// "errors"
// "fmt"
// "hash/crc32"
// "io"
// )
type Writer struct {
w io.Writer
}
// type Writer struct {
// w io.Writer
// }
func NewPNGWriter(w io.Writer) (*Writer, error) {
if _, err := io.WriteString(w, writeHeader); err != nil {
return nil, err
}
return &Writer{w}, nil
}
// func NewPNGWriter(w io.Writer) (*Writer, error) {
// if _, err := io.WriteString(w, writeHeader); err != nil {
// return nil, err
// }
// return &Writer{w}, nil
// }
func (w *Writer) WriteChunk(length int32, typ string, r io.Reader) error {
if err := binary.Write(w.w, binary.BigEndian, length); err != nil {
return err
}
if _, err := w.w.Write([]byte(typ)); err != nil {
return err
}
checksummer := crc32.NewIEEE()
checksummer.Write([]byte(typ))
if _, err := io.CopyN(io.MultiWriter(w.w, checksummer), r, int64(length)); err != nil {
return err
}
if err := binary.Write(w.w, binary.BigEndian, checksummer.Sum32()); err != nil {
return err
}
return nil
}
// func (w *Writer) WriteChunk(length int32, typ string, r io.Reader) error {
// if err := binary.Write(w.w, binary.BigEndian, length); err != nil {
// return err
// }
// if _, err := w.w.Write([]byte(typ)); err != nil {
// return err
// }
// checksummer := crc32.NewIEEE()
// checksummer.Write([]byte(typ))
// if _, err := io.CopyN(io.MultiWriter(w.w, checksummer), r, int64(length)); err != nil {
// return err
// }
// if err := binary.Write(w.w, binary.BigEndian, checksummer.Sum32()); err != nil {
// return err
// }
// return nil
// }
func WriteToPng(c *models.CharCardSpec, fpath, outfile string) error {
data, err := os.ReadFile(fpath)
if err != nil {
return err
}
jsonData, err := json.Marshal(c)
if err != nil {
return err
}
// Base64 encode the JSON data
base64Data := base64.StdEncoding.EncodeToString(jsonData)
pe := PngEmbed{
Key: cKey,
Value: base64Data,
}
w, err := WritetEXtToPngBytes(data, pe)
if err != nil {
return err
}
return os.WriteFile(outfile, w.Bytes(), 0666)
}
// func WWriteToPngriteToPng(c *models.CharCardSpec, fpath, outfile string) error {
// data, err := os.ReadFile(fpath)
// if err != nil {
// return err
// }
// jsonData, err := json.Marshal(c)
// if err != nil {
// return err
// }
// // Base64 encode the JSON data
// base64Data := base64.StdEncoding.EncodeToString(jsonData)
// pe := PngEmbed{
// Key: cKey,
// Value: base64Data,
// }
// w, err := WritetEXtToPngBytes(data, pe)
// if err != nil {
// return err
// }
// return os.WriteFile(outfile, w.Bytes(), 0666)
// }
func WritetEXtToPngBytes(inputBytes []byte, pe PngEmbed) (outputBytes bytes.Buffer, err error) {
if !(string(inputBytes[:8]) == header) {
return outputBytes, errors.New("wrong file format")
}
reader := bytes.NewReader(inputBytes)
pngr, err := NewPNGStepReader(reader)
if err != nil {
return outputBytes, fmt.Errorf("NewReader(): %s", err)
}
pngw, err := NewPNGWriter(&outputBytes)
if err != nil {
return outputBytes, fmt.Errorf("NewWriter(): %s", err)
}
for {
chunk, err := pngr.Next()
if err != nil {
if errors.Is(err, io.EOF) {
break
}
return outputBytes, fmt.Errorf("NextChunk(): %s", err)
}
if chunk.Type() != embType {
// IENDChunkType will only appear on the final iteration of a valid PNG
if chunk.Type() == IEND {
// This is where we inject tEXtChunkType as the penultimate chunk with the new value
newtEXtChunk := []byte(fmt.Sprintf(tEXtChunkDataSpecification, pe.Key, pe.Value))
if err := pngw.WriteChunk(int32(len(newtEXtChunk)), embType, bytes.NewBuffer(newtEXtChunk)); err != nil {
return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
}
// Now we end the buffer with IENDChunkType chunk
if err := pngw.WriteChunk(chunk.length, chunk.Type(), chunk); err != nil {
return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
}
} else {
// writes back original chunk to buffer
if err := pngw.WriteChunk(chunk.length, chunk.Type(), chunk); err != nil {
return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
}
}
} else {
if _, err := io.Copy(io.Discard, chunk); err != nil {
return outputBytes, fmt.Errorf("io.Copy(io.Discard, chunk): %s", err)
}
}
if err := chunk.Close(); err != nil {
return outputBytes, fmt.Errorf("chunk.Close(): %s", err)
}
}
return outputBytes, nil
}
// func WritetEXtToPngBytes(inputBytes []byte, pe PngEmbed) (outputBytes bytes.Buffer, err error) {
// if !(string(inputBytes[:8]) == header) {
// return outputBytes, errors.New("wrong file format")
// }
// reader := bytes.NewReader(inputBytes)
// pngr, err := NewPNGStepReader(reader)
// if err != nil {
// return outputBytes, fmt.Errorf("NewReader(): %s", err)
// }
// pngw, err := NewPNGWriter(&outputBytes)
// if err != nil {
// return outputBytes, fmt.Errorf("NewWriter(): %s", err)
// }
// for {
// chunk, err := pngr.Next()
// if err != nil {
// if errors.Is(err, io.EOF) {
// break
// }
// return outputBytes, fmt.Errorf("NextChunk(): %s", err)
// }
// if chunk.Type() != embType {
// // IENDChunkType will only appear on the final iteration of a valid PNG
// if chunk.Type() == IEND {
// // This is where we inject tEXtChunkType as the penultimate chunk with the new value
// newtEXtChunk := []byte(fmt.Sprintf(tEXtChunkDataSpecification, pe.Key, pe.Value))
// if err := pngw.WriteChunk(int32(len(newtEXtChunk)), embType, bytes.NewBuffer(newtEXtChunk)); err != nil {
// return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
// }
// // Now we end the buffer with IENDChunkType chunk
// if err := pngw.WriteChunk(chunk.length, chunk.Type(), chunk); err != nil {
// return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
// }
// } else {
// // writes back original chunk to buffer
// if err := pngw.WriteChunk(chunk.length, chunk.Type(), chunk); err != nil {
// return outputBytes, fmt.Errorf("WriteChunk(): %s", err)
// }
// }
// } else {
// if _, err := io.Copy(io.Discard, chunk); err != nil {
// return outputBytes, fmt.Errorf("io.Copy(io.Discard, chunk): %s", err)
// }
// }
// if err := chunk.Close(); err != nil {
// return outputBytes, fmt.Errorf("chunk.Close(): %s", err)
// }
// }
// return outputBytes, nil
// }