diff --git a/src/compiler/Write.go b/src/compiler/Write.go index cc1981f..f904b6b 100644 --- a/src/compiler/Write.go +++ b/src/compiler/Write.go @@ -1,7 +1,6 @@ package compiler import ( - "bufio" "io" "git.urbach.dev/cli/q/src/config" @@ -13,24 +12,20 @@ import ( ) // Write writes the executable to the given writer. -func (r *Result) Write(writer io.Writer) error { - return write(writer, r.Code, r.Data, r.DLLs) +func (r *Result) Write(writer io.WriteSeeker) { + write(writer, r.Code, r.Data, r.DLLs) } // write writes an executable file to the given writer. -func write(writer io.Writer, code []byte, data []byte, dlls dll.List) error { - buffer := bufio.NewWriter(writer) - +func write(writer io.WriteSeeker, code []byte, data []byte, dlls dll.List) { switch config.TargetOS { case config.Linux: - elf.Write(buffer, code, data) + elf.Write(writer, code, data) case config.Mac: - macho.Write(buffer, code, data) + macho.Write(writer, code, data) case config.Web: - wasm.Write(buffer, code, data) + wasm.Write(writer, code, data) case config.Windows: - pe.Write(buffer, code, data, dlls) + pe.Write(writer, code, data, dlls) } - - return buffer.Flush() } diff --git a/src/compiler/WriteFile.go b/src/compiler/WriteFile.go index 009562e..067ebfd 100644 --- a/src/compiler/WriteFile.go +++ b/src/compiler/WriteFile.go @@ -10,13 +10,7 @@ func (r *Result) WriteFile(path string) error { return err } - err = r.Write(file) - - if err != nil { - file.Close() - return err - } - + r.Write(file) err = file.Close() if err != nil { diff --git a/src/elf/ELF.go b/src/elf/ELF.go index 29c828a..0a03120 100644 --- a/src/elf/ELF.go +++ b/src/elf/ELF.go @@ -1,7 +1,6 @@ package elf import ( - "bytes" "encoding/binary" "io" @@ -21,7 +20,7 @@ type ELF struct { } // Write writes the ELF64 format to the given writer. -func Write(writer io.Writer, codeBytes []byte, dataBytes []byte) { +func Write(writer io.WriteSeeker, codeBytes []byte, dataBytes []byte) { sections := exe.MakeSections(HeaderEnd, codeBytes, dataBytes) code := sections[0] data := sections[1] @@ -76,9 +75,9 @@ func Write(writer io.Writer, codeBytes []byte, dataBytes []byte) { binary.Write(writer, binary.LittleEndian, &elf.Header) binary.Write(writer, binary.LittleEndian, &elf.CodeHeader) binary.Write(writer, binary.LittleEndian, &elf.DataHeader) - writer.Write(bytes.Repeat([]byte{0x00}, code.Padding)) + writer.Seek(int64(code.Padding), io.SeekCurrent) writer.Write(code.Bytes) - writer.Write(bytes.Repeat([]byte{0x00}, data.Padding)) + writer.Seek(int64(data.Padding), io.SeekCurrent) writer.Write(data.Bytes) if config.Sections { diff --git a/src/elf/ELF_test.go b/src/elf/ELF_test.go index 74bd71c..528bd75 100644 --- a/src/elf/ELF_test.go +++ b/src/elf/ELF_test.go @@ -1,12 +1,12 @@ package elf_test import ( - "io" "testing" "git.urbach.dev/cli/q/src/elf" + "git.urbach.dev/cli/q/src/test" ) func TestWrite(t *testing.T) { - elf.Write(io.Discard, nil, nil) + elf.Write(&test.Discard{}, nil, nil) } diff --git a/src/macho/MachO.go b/src/macho/MachO.go index b305b6a..8588b3d 100644 --- a/src/macho/MachO.go +++ b/src/macho/MachO.go @@ -1,7 +1,6 @@ package macho import ( - "bytes" "encoding/binary" "io" @@ -24,7 +23,7 @@ type MachO struct { } // Write writes the Mach-O format to the given writer. -func Write(writer io.Writer, codeBytes []byte, dataBytes []byte) { +func Write(writer io.WriteSeeker, codeBytes []byte, dataBytes []byte) { sections := exe.MakeSections(HeaderEnd, codeBytes, dataBytes) code := sections[0] data := sections[1] @@ -116,8 +115,8 @@ func Write(writer io.Writer, codeBytes []byte, dataBytes []byte) { binary.Write(writer, binary.LittleEndian, &m.CodeHeader) binary.Write(writer, binary.LittleEndian, &m.DataHeader) binary.Write(writer, binary.LittleEndian, &m.UnixThread) - writer.Write(bytes.Repeat([]byte{0x00}, code.Padding)) + writer.Seek(int64(code.Padding), io.SeekCurrent) writer.Write(code.Bytes) - writer.Write(bytes.Repeat([]byte{0x00}, data.Padding)) + writer.Seek(int64(data.Padding), io.SeekCurrent) writer.Write(data.Bytes) } diff --git a/src/macho/MachO_test.go b/src/macho/MachO_test.go index df2d04c..f3cf5f9 100644 --- a/src/macho/MachO_test.go +++ b/src/macho/MachO_test.go @@ -1,12 +1,12 @@ package macho_test import ( - "io" "testing" "git.urbach.dev/cli/q/src/macho" + "git.urbach.dev/cli/q/src/test" ) func TestWrite(t *testing.T) { - macho.Write(io.Discard, nil, nil) + macho.Write(&test.Discard{}, nil, nil) } diff --git a/src/pe/EXE.go b/src/pe/EXE.go index 740d0a5..210208b 100644 --- a/src/pe/EXE.go +++ b/src/pe/EXE.go @@ -24,7 +24,7 @@ type EXE struct { } // Write writes the EXE file to the given writer. -func Write(writer io.Writer, codeBytes []byte, dataBytes []byte, dlls dll.List) { +func Write(writer io.WriteSeeker, codeBytes []byte, dataBytes []byte, dlls dll.List) { var ( sections = exe.MakeSections(HeaderEnd, codeBytes, dataBytes, nil) code = sections[0] @@ -144,10 +144,10 @@ func Write(writer io.Writer, codeBytes []byte, dataBytes []byte, dlls dll.List) binary.Write(writer, binary.LittleEndian, &pe.NTHeader) binary.Write(writer, binary.LittleEndian, &pe.OptionalHeader64) binary.Write(writer, binary.LittleEndian, &pe.Sections) - writer.Write(bytes.Repeat([]byte{0x00}, code.Padding)) + writer.Seek(int64(code.Padding), io.SeekCurrent) writer.Write(code.Bytes) - writer.Write(bytes.Repeat([]byte{0x00}, data.Padding)) + writer.Seek(int64(data.Padding), io.SeekCurrent) writer.Write(data.Bytes) - writer.Write(bytes.Repeat([]byte{0x00}, imports.Padding)) + writer.Seek(int64(imports.Padding), io.SeekCurrent) writer.Write(imports.Bytes) } diff --git a/src/pe/EXE_test.go b/src/pe/EXE_test.go index d69d46a..4059778 100644 --- a/src/pe/EXE_test.go +++ b/src/pe/EXE_test.go @@ -1,12 +1,12 @@ package pe_test import ( - "io" "testing" "git.urbach.dev/cli/q/src/pe" + "git.urbach.dev/cli/q/src/test" ) func TestWrite(t *testing.T) { - pe.Write(io.Discard, nil, nil, nil) + pe.Write(&test.Discard{}, nil, nil, nil) } diff --git a/src/readme.md b/src/readme.md index 83641ce..417d18c 100644 --- a/src/readme.md +++ b/src/readme.md @@ -26,6 +26,7 @@ - [scope](scope) - Defines a `Scope` used for code blocks - [set](set) - Generic set implementation - [sizeof](sizeof) - Calculates the byte size of numbers +- [test](test) - Testing utilities - [token](token) - Converts a file to tokens with the `Tokenize` function - [types](types) - Type system - [x86](x86) - x86-64 implementation diff --git a/src/test/Discard.go b/src/test/Discard.go new file mode 100644 index 0000000..62cf0c0 --- /dev/null +++ b/src/test/Discard.go @@ -0,0 +1,7 @@ +package test + +// Discard implements a no-op WriteSeeker. +type Discard struct{} + +func (w *Discard) Write(_ []byte) (int, error) { return 0, nil } +func (w *Discard) Seek(_ int64, _ int) (int64, error) { return 0, nil }