1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400

// Package varlinklib provides shared utilities for working with varlink services,
// including high-level client wrappers with automatic retry logic and low-level
// socket connection helpers.
package varlinklib

import (
	"bufio"
	"context"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"log"
	"maps"
	"net"
	"net/http"
	"net/url"
	"strings"
	"sync"

	"github.com/varlink/go/varlink"
	"github.com/varlink/go/varlink/idl"
)

// =============================================================================
// High-level VarlinkClient with automatic retry
// =============================================================================

// VarlinkClient wraps a varlink connection and provides method introspection
// with automatic reconnection and retry on transport errors.
type VarlinkClient struct {
	conn      *varlink.Connection
	idl       *idl.IDL
	ifaceName string
	address   string
	mu        sync.Mutex
}

// NewVarlinkClient connects to a varlink service and introspects its interface
func NewVarlinkClient(ctx context.Context, address, interfaceName string) (*VarlinkClient, error) {
	// Connect to varlink service
	conn, err := varlink.NewConnection(ctx, address)
	if err != nil {
		return nil, fmt.Errorf("failed to connect to %s: %w", address, err)
	}

	// Get interface description
	description, err := conn.GetInterfaceDescription(ctx, interfaceName)
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to get interface description for %s: %w", interfaceName, err)
	}

	// Parse interface definition with IDL parser
	parsedIDL, err := idl.New(description)
	if err != nil {
		conn.Close()
		return nil, fmt.Errorf("failed to parse interface definition: %w", err)
	}

	return &VarlinkClient{
		conn:      conn,
		idl:       parsedIDL,
		ifaceName: interfaceName,
		address:   address,
	}, nil
}

// GetIDL returns the parsed varlink interface IDL
func (vc *VarlinkClient) GetIDL() *idl.IDL {
	return vc.idl
}

// reconnect closes the current connection and establishes a new one,
// re-introspecting the interface. This method is thread-safe.
func (vc *VarlinkClient) reconnect(ctx context.Context) error {
	vc.mu.Lock()
	defer vc.mu.Unlock()

	log.Printf("Attempting to reconnect to %s...\n", vc.address)

	// Close old connection if it exists
	if vc.conn != nil {
		vc.conn.Close()
	}

	// Connect to varlink service
	conn, err := varlink.NewConnection(ctx, vc.address)
	if err != nil {
		return fmt.Errorf("failed to reconnect to %s: %w", vc.address, err)
	}

	// Get interface description
	description, err := conn.GetInterfaceDescription(ctx, vc.ifaceName)
	if err != nil {
		conn.Close()
		return fmt.Errorf("failed to get interface description for %s: %w", vc.ifaceName, err)
	}

	// Parse interface definition with IDL parser
	parsedIDL, err := idl.New(description)
	if err != nil {
		conn.Close()
		return fmt.Errorf("failed to parse interface definition: %w", err)
	}

	// Update connection and IDL
	vc.conn = conn
	vc.idl = parsedIDL

	log.Printf("Successfully reconnected to %s\n", vc.address)
	return nil
}

// Call invokes a varlink method and returns the result.
// If a transport error occurs, it attempts to reconnect once and retry the call.
func (vc *VarlinkClient) Call(ctx context.Context, methodName string, params map[string]any) (map[string]any, error) {
	// Construct full method name: interface.method
	fullMethod := vc.ifaceName + "." + methodName

	// Make the call
	var result map[string]any
	err := vc.conn.Call(ctx, fullMethod, params, &result)

	// If successful, return immediately
	if err == nil {
		return result, nil
	}

	// If it's an application error (method not found, invalid params, etc.),
	// return immediately without reconnecting
	if IsVarlinkApplicationError(err) {
		return nil, fmt.Errorf("varlink call failed: %w", err)
	}

	// Transport error detected - log and attempt reconnection
	log.Printf("Transport error during varlink call: %v\n", err)

	if reconnectErr := vc.reconnect(ctx); reconnectErr != nil {
		return nil, fmt.Errorf("reconnection failed: %w (original error: %v)", reconnectErr, err)
	}

	// Retry the call after successful reconnection
	log.Printf("Retrying call to %s after reconnection\n", fullMethod)
	err = vc.conn.Call(ctx, fullMethod, params, &result)
	if err != nil {
		return nil, fmt.Errorf("varlink call failed after reconnection: %w", err)
	}

	return result, nil
}

// Close closes the varlink connection
func (vc *VarlinkClient) Close() error {
	if vc.conn != nil {
		return vc.conn.Close()
	}
	return nil
}

// FetchResources fetches multiple binary resources over a single HTTP-upgraded connection.
// This is used when a varlink method returns resource_type: "http_listing".
//
// Parameters:
//   - ctx: Context for cancellation
//   - methodName: Varlink method name (e.g., "GetMessage")
//   - params: Method parameters (must NOT include upgrade_to_http_serve - we add it)
//   - resourcePaths: List of HTTP paths to fetch (e.g., ["/parts/0", "/parts/1"])
//
// Returns:
//   - Map of resourcePath -> binary data
//   - Error if upgrade or fetch fails
//
// The function:
// 1. Opens a raw socket connection
// 2. Sends varlink call with upgrade_to_http_serve=true
// 3. Verifies upgrade was successful
// 4. Sends HTTP GET for each resource path sequentially over same connection
// 5. Returns all fetched data
func (vc *VarlinkClient) FetchResources(ctx context.Context, methodName string, params map[string]any, resourcePaths []string) (map[string][]byte, error) {
	// Create a copy of params and add upgrade flag
	upgradeParams := make(map[string]any)
	maps.Copy(upgradeParams, params)
	upgradeParams["upgrade_to_http_serve"] = true

	// Dial raw socket
	conn, err := DialVarlinkSocket(vc.address)
	if err != nil {
		return nil, fmt.Errorf("failed to dial socket: %w", err)
	}
	defer conn.Close()

	// Build and send varlink call
	fullMethod := vc.ifaceName + "." + methodName
	call := map[string]any{
		"method":     fullMethod,
		"parameters": upgradeParams,
	}

	callJSON, err := json.Marshal(call)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal call: %w", err)
	}

	// Send with null terminator
	if _, err := conn.Write(append(callJSON, 0)); err != nil {
		return nil, fmt.Errorf("failed to write call: %w", err)
	}

	// Read upgrade confirmation
	scanner := bufio.NewScanner(conn)
	scanner.Split(scanNullTerminated)
	if !scanner.Scan() {
		if err := scanner.Err(); err != nil {
			return nil, fmt.Errorf("failed to read upgrade response: %w", err)
		}
		return nil, fmt.Errorf("no upgrade response from service")
	}

	var upgradeReply map[string]any
	if err := json.Unmarshal(scanner.Bytes(), &upgradeReply); err != nil {
		return nil, fmt.Errorf("failed to parse upgrade response: %w", err)
	}

	// Check for error
	if errMsg, hasErr := upgradeReply["error"]; hasErr {
		return nil, fmt.Errorf("varlink error during upgrade: %v", errMsg)
	}

	// Verify upgrade happened
	var upgraded bool
	if params, ok := upgradeReply["parameters"].(map[string]any); ok {
		upgraded, _ = params["upgraded"].(bool)
	}
	if !upgraded {
		return nil, fmt.Errorf("service did not confirm upgrade")
	}

	// Now connection is in HTTP mode - fetch all resources
	results := make(map[string][]byte)

	for _, resourcePath := range resourcePaths {
		// Send HTTP GET request
		httpReq := fmt.Sprintf("GET %s HTTP/1.1\r\nHost: localhost\r\nConnection: keep-alive\r\n\r\n", resourcePath)
		if _, err := conn.Write([]byte(httpReq)); err != nil {
			return nil, fmt.Errorf("failed to send HTTP request for %s: %w", resourcePath, err)
		}

		// Read HTTP response
		httpResp, err := http.ReadResponse(bufio.NewReader(conn), nil)
		if err != nil {
			return nil, fmt.Errorf("failed to read HTTP response for %s: %w", resourcePath, err)
		}

		if httpResp.StatusCode != http.StatusOK {
			httpResp.Body.Close()
			return nil, fmt.Errorf("HTTP request for %s returned status %d", resourcePath, httpResp.StatusCode)
		}

		// Read body
		body, err := io.ReadAll(httpResp.Body)
		httpResp.Body.Close()
		if err != nil {
			return nil, fmt.Errorf("failed to read body for %s: %w", resourcePath, err)
		}

		results[resourcePath] = body
	}

	return results, nil
}

// =============================================================================
// Low-level socket helpers for raw varlink protocol
// =============================================================================

// scanNullTerminated is a split function for bufio.Scanner that splits on \0 bytes
func scanNullTerminated(data []byte, atEOF bool) (advance int, token []byte, err error) {
	// Look for null byte
	for i := range data {
		if data[i] == 0 {
			// Found null terminator, return everything before it
			return i + 1, data[0:i], nil
		}
	}

	// If we're at EOF and have data, return it
	if atEOF && len(data) > 0 {
		return len(data), data, nil
	}

	// Need more data
	return 0, nil, nil
}

// DialVarlinkSocket parses a varlink address and dials the socket.
// Supports unix:// and tcp:// schemes.
func DialVarlinkSocket(address string) (net.Conn, error) {
	u, err := url.Parse(address)
	if err != nil {
		return nil, fmt.Errorf("invalid address: %w", err)
	}

	switch u.Scheme {
	case "unix":
		return net.Dial("unix", u.Path)
	case "tcp":
		return net.Dial("tcp", u.Host)
	default:
		return nil, fmt.Errorf("unsupported scheme: %s (expected unix or tcp)", u.Scheme)
	}
}

// DialVarlinkSocketWithRetry attempts to dial a varlink socket with retry logic.
// If the initial dial fails with a transport error, it retries once.
func DialVarlinkSocketWithRetry(address string) (net.Conn, error) {
	conn, err := DialVarlinkSocket(address)
	if err == nil {
		return conn, nil
	}

	// If it's a transport error, retry once
	if IsTransportError(err) {
		log.Printf("Transport error during dial to %s: %v, retrying...\n", address, err)
		conn, retryErr := DialVarlinkSocket(address)
		if retryErr != nil {
			return nil, fmt.Errorf("dial failed after retry: %w (original error: %v)", retryErr, err)
		}
		log.Printf("Successfully connected to %s after retry\n", address)
		return conn, nil
	}

	// Not a transport error, return immediately
	return nil, err
}

// =============================================================================
// Error classification utilities
// =============================================================================

// IsVarlinkApplicationError checks if an error is a varlink application error
// (as opposed to a transport/connection error). Application errors should not
// trigger reconnection.
func IsVarlinkApplicationError(err error) bool {
	if err == nil {
		return false
	}

	// Check if it's a varlink.Error (application-level error from the service)
	var varlinkErr *varlink.Error
	if errors.As(err, &varlinkErr) {
		return true
	}

	// Check for specific varlink service error types
	var methodNotFound *varlink.MethodNotFound
	var methodNotImplemented *varlink.MethodNotImplemented
	var invalidParameter *varlink.InvalidParameter
	var interfaceNotFound *varlink.InterfaceNotFound

	if errors.As(err, &methodNotFound) ||
		errors.As(err, &methodNotImplemented) ||
		errors.As(err, &invalidParameter) ||
		errors.As(err, &interfaceNotFound) {
		return true
	}

	return false
}

// IsTransportError checks if an error is a transport/connection error
// (as opposed to an application error). Transport errors should trigger retry.
func IsTransportError(err error) bool {
	if err == nil {
		return false
	}

	// Check for common transport errors
	if err == io.EOF || err == io.ErrClosedPipe || err == io.ErrUnexpectedEOF {
		return true
	}

	// Check for network errors using errors.As
	var netErr *net.OpError
	var urlErr *url.Error
	if errors.As(err, &netErr) || errors.As(err, &urlErr) {
		return true
	}

	// Check error message for common patterns
	errMsg := err.Error()
	if strings.Contains(errMsg, "broken pipe") ||
		strings.Contains(errMsg, "connection reset") ||
		strings.Contains(errMsg, "connection refused") ||
		strings.Contains(errMsg, "no such file or directory") { // unix socket doesn't exist
		return true
	}

	return false
}