diff --git a/main.go b/main.go index 591afb8..92e984a 100644 --- a/main.go +++ b/main.go @@ -1,18 +1,22 @@ package main import ( + "bufio" + "encoding/hex" "flag" "fmt" "image" "image/color" _ "image/gif" _ "image/jpeg" - _ "image/png" + "image/png" "log" "math/rand" "net" + "net/textproto" "os" "runtime/pprof" + "strings" "time" ) @@ -25,6 +29,7 @@ var connections = flag.Int("connections", 4, "Number of simultaneous connections var address = flag.String("host", "127.0.0.1:1337", "Server address") var runtime = flag.String("runtime", "60s", "exit after timeout") var shuffle = flag.Bool("shuffle", false, "pixel send ordering") +var fetchImgPath = flag.String("fetch-image", "", "path to save the fetched pixel state to") func main() { flag.Parse() @@ -52,15 +57,34 @@ func main() { offset := image.Pt(*image_offsetx, *image_offsety) img := readImage(*image_path) + fetchedImg := image.NewNRGBA(img.Bounds().Add(offset)) + + if *fetchImgPath != "" { + // using img.SubImage to distribute tasks is nice, as we can also parallelize command generation easily! + // @cleanup: use a box tiling algo instead of hardcoding + b := fetchedImg.Bounds() + b1 := image.Rectangle{b.Min, b.Size().Div(2)} + b2 := b1.Add(image.Pt(b1.Dx(), 0)) + b3 := b1.Add(image.Pt(0, b1.Dy())) + b4 := b1.Add(b1.Size()) + go fetchImage(fetchedImg.SubImage(b1).(*image.NRGBA)) + go fetchImage(fetchedImg.SubImage(b2).(*image.NRGBA)) + go fetchImage(fetchedImg.SubImage(b3).(*image.NRGBA)) + go fetchImage(fetchedImg.SubImage(b4).(*image.NRGBA)) + *connections -= 4 + } + // Generate and split messages into equal chunks commands := genCommands(img, offset) if *shuffle { shuffleCommands(commands) } - commandGroups := chunkCommands(commands, *connections) - for _, messages := range commandGroups { - go bomb(messages) + if *connections > 0 { + commandGroups := chunkCommands(commands, *connections) + for _, messages := range commandGroups { + go bomb(messages) + } } // Terminate after timeout to save resources @@ -69,6 +93,10 @@ func main() { log.Fatal("Invalid runtime specified: " + err.Error()) } time.Sleep(timer) + + if *fetchImgPath != "" { + writeImage(*fetchImgPath, fetchedImg) + } } func bomb(messages []byte) { @@ -102,6 +130,17 @@ func readImage(path string) (img image.Image) { return img } +func writeImage(path string, img image.Image) { + f, err := os.Create(path) + if err != nil { + log.Fatal(err) + } + if err := png.Encode(f, img); err != nil { + f.Close() + log.Fatal(err) + } +} + // Creates message based on given image func genCommands(img image.Image, offset image.Point) (commands [][]byte) { b := img.Bounds() @@ -149,3 +188,47 @@ func shuffleCommands(slice [][]byte) { slice[i], slice[j] = slice[j], slice[i] } } + +func fetchImage(img *image.NRGBA) { + // FIXME @speed: this is unusably s l o w w w + // bottleneck seems to be our pixel reading/parsing code. cpuprofile! + // -> should buffer it just as in bomb() + + conn, err := net.Dial("tcp", *address) + if err != nil { + log.Fatal(err) + } + defer conn.Close() + + reader := bufio.NewReader(conn) + tp := textproto.NewReader(reader) + + b := img.Bounds() + for { + for x := b.Min.X; x < b.Max.X; x++ { + for y := b.Min.Y; y < b.Max.Y; y++ { + // request pixel + fmt.Fprintf(conn, "PX %d %d\n", x, y) + if err != nil { + log.Fatal(err) + } + + // read pixel + // @speed try to run this in a separate goroutine? + // we probably need to buffer the responses then + res, err := tp.ReadLine() + if err != nil { + log.Fatal(err) + } + res2 := strings.Split(res, " ") + col, _ := hex.DecodeString(res2[3]) + img.Set(x, y, color.NRGBA{ + uint8(col[0]), + uint8(col[1]), + uint8(col[2]), + 255, + }) + } + } + } +}