WeKnora deployment source with the reviewed wiki identity, summary and recovery fixes

docreader/client/auth.go

@@ -0,0 +1,134 @@
+// Package client provides a docreader gRPC client and the shared TLS / token
+// authentication helpers used by both the standalone Go SDK in this package
+// and the internal docparser wrapper. Keep all auth/TLS construction here so
+// the two call sites cannot drift on security defaults.
+package client
+
+import (
+	"context"
+	"crypto/tls"
+	"crypto/x509"
+	"fmt"
+	"os"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/credentials/insecure"
+)
+
+// AuthConfig holds the docreader gRPC client TLS / token configuration.
+type AuthConfig struct {
+	TLSEnabled bool
+	CertFile   string
+	KeyFile    string
+	CAFile     string
+	// ServerName overrides the SNI / certificate-host check on the client.
+	// When empty, the address passed to Dial is used by Go's TLS stack.
+	ServerName string
+
+	AuthToken string
+}
+
+// LoadAuthConfigFromEnv reads docreader gRPC auth knobs from the process
+// environment. The caller must pass the result to BuildDialOptions to apply
+// them to a gRPC connection.
+func LoadAuthConfigFromEnv() *AuthConfig {
+	return &AuthConfig{
+		TLSEnabled: os.Getenv("GRPC_TLS_ENABLED") == "true",
+		CertFile:   os.Getenv("GRPC_TLS_CERT"),
+		KeyFile:    os.Getenv("GRPC_TLS_KEY"),
+		CAFile:     os.Getenv("GRPC_TLS_CA"),
+		ServerName: os.Getenv("GRPC_TLS_SERVER_NAME"),
+		AuthToken:  os.Getenv("GRPC_AUTH_TOKEN"),
+	}
+}
+
+// BuildDialOptions returns the gRPC DialOptions that apply the configured
+// transport credentials and per-RPC token. Callers should append their own
+// per-call options (load balancer, message size, etc.).
+func (c *AuthConfig) BuildDialOptions(maxMsgSize int) ([]grpc.DialOption, error) {
+	opts := []grpc.DialOption{
+		grpc.WithDefaultServiceConfig(`{"loadBalancingPolicy":"round_robin"}`),
+		grpc.WithDefaultCallOptions(
+			grpc.MaxCallRecvMsgSize(maxMsgSize),
+			grpc.MaxCallSendMsgSize(maxMsgSize),
+		),
+	}
+
+	if c.TLSEnabled {
+		creds, err := c.buildTLSCredentials()
+		if err != nil {
+			return nil, fmt.Errorf("failed to build TLS credentials: %w", err)
+		}
+		opts = append(opts, grpc.WithTransportCredentials(creds))
+		Logger.Printf("INFO: TLS enabled for gRPC client")
+	} else {
+		opts = append(opts, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	}
+
+	if c.AuthToken != "" {
+		// Only allow per-RPC tokens to ride a secured channel. This mirrors
+		// gRPC's own oauth2 credentials behaviour and prevents the bearer
+		// token from leaking on plaintext connections.
+		opts = append(opts, grpc.WithPerRPCCredentials(&tokenAuth{
+			token:           c.AuthToken,
+			requireTLSGuard: c.TLSEnabled,
+		}))
+		Logger.Printf("INFO: Token authentication enabled for gRPC client (TLS=%v)", c.TLSEnabled)
+	}
+
+	return opts, nil
+}
+
+func (c *AuthConfig) buildTLSCredentials() (credentials.TransportCredentials, error) {
+	tlsConfig := &tls.Config{
+		MinVersion: tls.VersionTLS12,
+		ServerName: c.ServerName,
+	}
+
+	if c.CAFile != "" {
+		caCert, err := os.ReadFile(c.CAFile)
+		if err != nil {
+			return nil, fmt.Errorf("failed to read CA certificate: %w", err)
+		}
+		certPool := x509.NewCertPool()
+		if !certPool.AppendCertsFromPEM(caCert) {
+			return nil, fmt.Errorf("failed to parse CA certificate")
+		}
+		tlsConfig.RootCAs = certPool
+	}
+
+	switch {
+	case c.CertFile != "" && c.KeyFile != "":
+		cert, err := tls.LoadX509KeyPair(c.CertFile, c.KeyFile)
+		if err != nil {
+			return nil, fmt.Errorf("failed to load client certificate: %w", err)
+		}
+		tlsConfig.Certificates = []tls.Certificate{cert}
+		Logger.Printf("INFO: mTLS enabled (client certificate loaded)")
+	case c.CertFile != "" || c.KeyFile != "":
+		return nil, fmt.Errorf(
+			"GRPC_TLS_CERT and GRPC_TLS_KEY must be set together for mTLS",
+		)
+	}
+
+	return credentials.NewTLS(tlsConfig), nil
+}
+
+type tokenAuth struct {
+	token string
+	// requireTLSGuard mirrors AuthConfig.TLSEnabled; we expose it via
+	// RequireTransportSecurity so gRPC will refuse to send the bearer token
+	// over an insecure connection when the operator has enabled TLS.
+	requireTLSGuard bool
+}
+
+func (t *tokenAuth) GetRequestMetadata(ctx context.Context, uri ...string) (map[string]string, error) {
+	return map[string]string{
+		"authorization": "Bearer " + t.token,
+	}, nil
+}
+
+func (t *tokenAuth) RequireTransportSecurity() bool {
+	return t.requireTLSGuard
+}

docreader/client/client.go

@@ -0,0 +1,113 @@
+package client
+
+import (
+	"fmt"
+	"log"
+	"os"
+	"strconv"
+	"time"
+
+	"github.com/Tencent/WeKnora/docreader/proto"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/resolver"
+)
+
+func getMaxMessageSize() int {
+	if sizeStr := os.Getenv("MAX_FILE_SIZE_MB"); sizeStr != "" {
+		if size, err := strconv.Atoi(sizeStr); err == nil && size > 0 {
+			if size > 20*1024 {
+				size = 20 * 1024
+			}
+			return size * 1024 * 1024
+		}
+	}
+	return 50 * 1024 * 1024
+}
+
+var Logger = log.New(os.Stdout, "[DocReader] ", log.LstdFlags|log.Lmicroseconds)
+
+// ImageRefInfo represents an image reference from a converted document.
+type ImageRefInfo struct {
+	Filename    string
+	OriginalRef string
+	MimeType    string
+	StorageKey  string
+}
+
+// Client represents a DocReader service client.
+type Client struct {
+	conn *grpc.ClientConn
+	proto.DocReaderClient
+	debug bool
+}
+
+func NewClient(addr string) (*Client, error) {
+	authConfig := LoadAuthConfigFromEnv()
+	return NewClientWithAuth(addr, authConfig)
+}
+
+func NewClientWithAuth(addr string, authConfig *AuthConfig) (*Client, error) {
+	Logger.Printf("INFO: Creating new DocReader client connecting to %s", addr)
+
+	maxMsgSize := getMaxMessageSize()
+
+	if authConfig == nil {
+		authConfig = &AuthConfig{}
+	}
+
+	opts, err := authConfig.BuildDialOptions(maxMsgSize)
+	if err != nil {
+		Logger.Printf("ERROR: Failed to build dial options: %v", err)
+		return nil, err
+	}
+
+	resolver.SetDefaultScheme("dns")
+
+	startTime := time.Now()
+	conn, err := grpc.Dial("dns:///"+addr, opts...)
+	if err != nil {
+		Logger.Printf("ERROR: Failed to connect to DocReader service: %v", err)
+		return nil, err
+	}
+	Logger.Printf("INFO: Successfully connected to DocReader service in %v", time.Since(startTime))
+
+	return &Client{
+		conn:            conn,
+		DocReaderClient: proto.NewDocReaderClient(conn),
+		debug:           false,
+	}, nil
+}
+
+func (c *Client) Close() error {
+	Logger.Printf("INFO: Closing DocReader client connection")
+	return c.conn.Close()
+}
+
+func (c *Client) SetDebug(debug bool) {
+	c.debug = debug
+}
+
+func (c *Client) Log(level string, format string, args ...interface{}) {
+	if level == "DEBUG" && !c.debug {
+		return
+	}
+	Logger.Printf("%s: %s", level, fmt.Sprintf(format, args...))
+}
+
+// GetImageRefsFromResponse extracts image references from a ReadResponse.
+func GetImageRefsFromResponse(resp *proto.ReadResponse) []ImageRefInfo {
+	if resp == nil || len(resp.ImageRefs) == 0 {
+		return nil
+	}
+
+	refs := make([]ImageRefInfo, 0, len(resp.ImageRefs))
+	for _, ref := range resp.ImageRefs {
+		refs = append(refs, ImageRefInfo{
+			Filename:    ref.Filename,
+			OriginalRef: ref.OriginalRef,
+			MimeType:    ref.MimeType,
+			StorageKey:  ref.StorageKey,
+		})
+	}
+	return refs
+}

docreader/client/client_test.go

@@ -0,0 +1,95 @@
+package client
+
+import (
+	"context"
+	"log"
+	"os"
+	"testing"
+	"time"
+
+	"github.com/Tencent/WeKnora/docreader/proto"
+)
+
+func init() {
+	log.SetOutput(os.Stdout)
+	log.SetFlags(log.LstdFlags | log.Lmicroseconds | log.Lshortfile)
+	log.Println("INFO: Initializing DocReader client tests")
+}
+
+func requireLiveDocReaderClient(t *testing.T) *Client {
+	t.Helper()
+
+	client, err := NewClient("localhost:50051")
+	if err != nil {
+		t.Fatalf("Failed to create client: %v", err)
+	}
+	t.Cleanup(func() { client.Close() })
+
+	ctx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
+	defer cancel()
+	if _, err := client.ListEngines(ctx, &proto.ListEnginesRequest{}); err != nil {
+		t.Skipf("DocReader gRPC server not available at localhost:50051: %v", err)
+	}
+
+	return client
+}
+
+func TestReadURL(t *testing.T) {
+	client := requireLiveDocReaderClient(t)
+	client.SetDebug(true)
+
+	startTime := time.Now()
+	resp, err := client.Read(
+		context.Background(),
+		&proto.ReadRequest{
+			Url:   "https://example.com",
+			Title: "test",
+		},
+	)
+	log.Printf("INFO: Read(URL) completed in %v", time.Since(startTime))
+
+	if err != nil {
+		t.Fatalf("Read failed: %v", err)
+	}
+	if resp.Error != "" {
+		t.Fatalf("Read returned error: %s", resp.Error)
+	}
+	if resp.MarkdownContent == "" {
+		t.Error("Expected non-empty markdown content")
+	}
+	log.Printf("INFO: content_len=%d, images=%d", len(resp.MarkdownContent), len(resp.ImageRefs))
+}
+
+func TestReadFile(t *testing.T) {
+	client := requireLiveDocReaderClient(t)
+	client.SetDebug(true)
+
+	fileContent, err := os.ReadFile("../testdata/test.md")
+	if err != nil {
+		t.Fatalf("Failed to read test file: %v", err)
+	}
+
+	startTime := time.Now()
+	resp, err := client.Read(
+		context.Background(),
+		&proto.ReadRequest{
+			FileContent: fileContent,
+			FileName:    "test.md",
+			FileType:    "md",
+		},
+	)
+	log.Printf("INFO: Read(file) completed in %v", time.Since(startTime))
+
+	if err != nil {
+		t.Fatalf("Read failed: %v", err)
+	}
+	if resp.Error != "" {
+		t.Fatalf("Read returned error: %s", resp.Error)
+	}
+	if resp.MarkdownContent == "" {
+		t.Error("Expected non-empty markdown content")
+	}
+
+	imageRefs := GetImageRefsFromResponse(resp)
+	log.Printf("INFO: content_len=%d, images=%d", len(resp.MarkdownContent), len(imageRefs))
+}