ds2api/internal/deepseek/pow.go

package deepseek

import (
	"context"
	"encoding/base64"
	"encoding/binary"
	"encoding/json"
	"errors"
	"math"
	"os"
	"sync"

	"ds2api/internal/config"

	"github.com/tetratelabs/wazero"
	"github.com/tetratelabs/wazero/api"
)

type PowSolver struct {
	wasmPath string
	once     sync.Once
	err      error

	runtime  wazero.Runtime
	compiled wazero.CompiledModule
}

func NewPowSolver(wasmPath string) *PowSolver {
	return &PowSolver{wasmPath: wasmPath}
}

func (p *PowSolver) init(ctx context.Context) error {
	p.once.Do(func() {
		wasmBytes, err := os.ReadFile(p.wasmPath)
		if err != nil {
			p.err = err
			return
		}
		p.runtime = wazero.NewRuntime(ctx)
		p.compiled, p.err = p.runtime.CompileModule(ctx, wasmBytes)
	})
	return p.err
}

func (p *PowSolver) Compute(ctx context.Context, challenge map[string]any) (int64, error) {
	if err := p.init(ctx); err != nil {
		return 0, err
	}
	algo, _ := challenge["algorithm"].(string)
	if algo != "DeepSeekHashV1" {
		return 0, errors.New("unsupported algorithm")
	}
	challengeStr, _ := challenge["challenge"].(string)
	salt, _ := challenge["salt"].(string)
	signature, _ := challenge["signature"].(string)
	targetPath, _ := challenge["target_path"].(string)
	_ = signature
	_ = targetPath

	difficulty := toFloat64(challenge["difficulty"], 144000)
	expireAt := toInt64(challenge["expire_at"], 1680000000)
	prefix := salt + "_" + itoa(expireAt) + "_"

	mod, err := p.runtime.InstantiateModule(ctx, p.compiled, wazero.NewModuleConfig())
	if err != nil {
		return 0, err
	}
	defer mod.Close(ctx)

	mem := mod.Memory()
	if mem == nil {
		return 0, errors.New("wasm memory missing")
	}
	stackFn := mod.ExportedFunction("__wbindgen_add_to_stack_pointer")
	allocFn := mod.ExportedFunction("__wbindgen_export_0")
	solveFn := mod.ExportedFunction("wasm_solve")
	if stackFn == nil || allocFn == nil || solveFn == nil {
		return 0, errors.New("required wasm exports missing")
	}

	retPtrs, err := stackFn.Call(ctx, uint64(uint32(^uint32(15)))) // -16 i32
	if err != nil || len(retPtrs) == 0 {
		return 0, errors.New("stack alloc failed")
	}
	retptr := uint32(retPtrs[0])
	defer stackFn.Call(ctx, 16)

	chPtr, chLen, err := writeUTF8(ctx, allocFn, mem, challengeStr)
	if err != nil {
		return 0, err
	}
	prefixPtr, prefixLen, err := writeUTF8(ctx, allocFn, mem, prefix)
	if err != nil {
		return 0, err
	}

	if _, err := solveFn.Call(ctx,
		uint64(retptr),
		uint64(chPtr), uint64(chLen),
		uint64(prefixPtr), uint64(prefixLen),
		math.Float64bits(difficulty),
	); err != nil {
		return 0, err
	}

	statusBytes, ok := mem.Read(retptr, 4)
	if !ok {
		return 0, errors.New("read status failed")
	}
	status := int32(binary.LittleEndian.Uint32(statusBytes))
	valueBytes, ok := mem.Read(retptr+8, 8)
	if !ok {
		return 0, errors.New("read value failed")
	}
	value := math.Float64frombits(binary.LittleEndian.Uint64(valueBytes))
	if status == 0 {
		return 0, errors.New("pow solve failed")
	}
	return int64(value), nil
}

func writeUTF8(ctx context.Context, allocFn api.Function, mem api.Memory, text string) (uint32, uint32, error) {
	data := []byte(text)
	res, err := allocFn.Call(ctx, uint64(len(data)), 1)
	if err != nil || len(res) == 0 {
		return 0, 0, errors.New("alloc failed")
	}
	ptr := uint32(res[0])
	if !mem.Write(ptr, data) {
		return 0, 0, errors.New("mem write failed")
	}
	return ptr, uint32(len(data)), nil
}

func BuildPowHeader(challenge map[string]any, answer int64) (string, error) {
	payload := map[string]any{
		"algorithm":   challenge["algorithm"],
		"challenge":   challenge["challenge"],
		"salt":        challenge["salt"],
		"answer":      answer,
		"signature":   challenge["signature"],
		"target_path": challenge["target_path"],
	}
	b, err := json.Marshal(payload)
	if err != nil {
		return "", err
	}
	return base64.StdEncoding.EncodeToString(b), nil
}

func toFloat64(v any, d float64) float64 {
	switch n := v.(type) {
	case float64:
		return n
	case int:
		return float64(n)
	case int64:
		return float64(n)
	default:
		return d
	}
}

func toInt64(v any, d int64) int64 {
	switch n := v.(type) {
	case float64:
		return int64(n)
	case int:
		return int64(n)
	case int64:
		return n
	default:
		return d
	}
}

func itoa(n int64) string {
	b, _ := json.Marshal(n)
	return string(b)
}

func PreloadWASM(wasmPath string) {
	solver := NewPowSolver(wasmPath)
	if err := solver.init(context.Background()); err != nil {
		config.Logger.Warn("[WASM] preload failed", "error", err)
		return
	}
	config.Logger.Info("[WASM] module preloaded", "path", wasmPath)
}