GH-46315 [C#] Apache Arrow Flight Middleware

csharp/src/Apache.Arrow.Flight.Sql/FlightSqlServer.cs

@@ -285,6 +285,7 @@ public override Task DoGet(FlightTicket ticket, FlightServerRecordBatchStreamWri
     public override Task DoAction(FlightAction action, IAsyncStreamWriter<FlightResult> responseStream, ServerCallContext context)
     {
         Logger?.LogTrace("Executing Flight SQL DoAction: {ActionType}", action.Type);
+
         switch (action.Type)
         {
             case SqlAction.CreateRequest:

csharp/src/Apache.Arrow.Flight/Middleware/CallHeaders.cs

@@ -0,0 +1,80 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using System;
+using System.Collections;
+using System.Collections.Generic;
+using System.Linq;
+using Apache.Arrow.Flight.Middleware.Interfaces;
+using Grpc.Core;
+
+namespace Apache.Arrow.Flight.Middleware;
+
+public class CallHeaders : ICallHeaders, IEnumerable<KeyValuePair<string, string>>
+{
+    private readonly Metadata _metadata;
+
+    public CallHeaders(Metadata metadata)
+    {
+        _metadata = metadata;
+    }
+
+    public void Add(string key, string value) => _metadata.Add(key, value);
+
+    public bool ContainsKey(string key) => _metadata.Any(h => KeyEquals(h.Key, key));
+
+    public IEnumerator<KeyValuePair<string, string>> GetEnumerator()
+    {
+        foreach (var entry in _metadata)
+            yield return new KeyValuePair<string, string>(entry.Key, entry.Value);
+    }
+
+    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+
+    public string this[string key]
+    {
+        get
+        {
+            var entry = _metadata.FirstOrDefault(h =>  KeyEquals(h.Key, key));
+            return entry?.Value;
+        }
+        set
+        {
+            var entry = _metadata.FirstOrDefault(h =>  KeyEquals(h.Key, key));
+            if (entry != null) _metadata.Remove(entry);
+            _metadata.Add(key, value);
+        }
+    }
+
+    public string Get(string key) => this[key];
+
+    public byte[] GetBytes(string key) =>
+        _metadata.FirstOrDefault(h => KeyEquals(h.Key, key))?.ValueBytes;
+
+    public IEnumerable<string> GetAll(string key) =>
+        _metadata.Where(h => KeyEquals(h.Key, key)).Select(h => h.Value);
+
+    public IEnumerable<byte[]> GetAllBytes(string key) =>
+        _metadata.Where(h => KeyEquals(h.Key, key)).Select(h => h.ValueBytes);
+
+    public void Insert(string key, string value) => Add(key, value);
+
+    public void Insert(string key, byte[] value) => _metadata.Add(key, value);
+
+    public ISet<string> Keys => new HashSet<string>(_metadata.Select(h => h.Key));
+
+    private static bool KeyEquals(string a, string b) =>
+        string.Equals(a, b, StringComparison.OrdinalIgnoreCase);
+}

csharp/src/Apache.Arrow.Flight/Middleware/CallInfo.cs

@@ -0,0 +1,35 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using Grpc.Core;
+
+namespace Apache.Arrow.Flight.Middleware;
+
+public readonly struct CallInfo
+{
+    public string Method { get; }
+    public MethodType MethodType { get; }
+
+    public CallInfo(string method, MethodType methodType)
+    {
+        Method = method;
+        MethodType = methodType;
+    }
+
+    public override string ToString()
+    {
+        return $"{MethodType}: {Method}";
+    }
+}

csharp/src/Apache.Arrow.Flight/Middleware/ClientCookieMiddleware.cs

@@ -0,0 +1,75 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using System.Collections.Generic;
+using Apache.Arrow.Flight.Middleware.Interfaces;
+using Grpc.Core;
+using Microsoft.Extensions.Logging;
+
+namespace Apache.Arrow.Flight.Middleware;
+
+public class ClientCookieMiddleware : IFlightClientMiddleware
+{
+    private readonly ClientCookieMiddlewareFactory _factory;
+    private readonly ILogger<ClientCookieMiddleware> _logger;
+    private const string SetCookieHeader = "Set-Cookie";
+    private const string CookieHeader = "Cookie";
+
+    public ClientCookieMiddleware(ClientCookieMiddlewareFactory factory,
+        ILogger<ClientCookieMiddleware> logger)
+    {
+        _factory = factory;
+        _logger = logger;
+    }
+
+    public void OnBeforeSendingHeaders(ICallHeaders outgoingHeaders)
+    {
+        if (_factory.Cookies.IsEmpty)
+            return;
+        var cookieValue = GetValidCookiesAsString();
+        if (!string.IsNullOrEmpty(cookieValue))
+        {
+            outgoingHeaders.Insert(CookieHeader, cookieValue);
+        }
+    }
+
+    public void OnHeadersReceived(ICallHeaders incomingHeaders)
+    {
+        var setCookies = incomingHeaders.GetAll(SetCookieHeader);
+        _factory.UpdateCookies(setCookies);
+    }
+
+    public void OnCallCompleted(Status status, Metadata trailers)
+    {
+        // ingest: status and/or metadata trailers
+    }
+
+    private string GetValidCookiesAsString()
+    {
+        var cookieList = new List<string>();
+        foreach (var entry in _factory.Cookies)
+        {
+            if (entry.Value.Expired)
+            {
+                _factory.Cookies.TryRemove(entry.Key, out _);
+            }
+            else
+            {
+                cookieList.Add(entry.Value.ToString());
+            }
+        }
+        return string.Join("; ", cookieList);
+    }
+}

csharp/src/Apache.Arrow.Flight/Middleware/ClientCookieMiddlewareFactory.cs

@@ -0,0 +1,67 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using System;
+using System.Collections.Concurrent;
+using System.Collections.Generic;
+using System.Globalization;
+using System.Net;
+using Apache.Arrow.Flight.Middleware.Extensions;
+using Apache.Arrow.Flight.Middleware.Interfaces;
+using Microsoft.Extensions.Logging;
+
+namespace Apache.Arrow.Flight.Middleware;
+
+public class ClientCookieMiddlewareFactory : IFlightClientMiddlewareFactory
+{
+    public readonly ConcurrentDictionary<string, Cookie> Cookies = new(StringComparer.OrdinalIgnoreCase);
+    private readonly ILogger<ClientCookieMiddleware> _logger;
+
+    public ClientCookieMiddlewareFactory(ILoggerFactory loggerFactory)
+    {
+        _logger = loggerFactory.CreateLogger<ClientCookieMiddleware>();
+    }
+
+    public IFlightClientMiddleware OnCallStarted(CallInfo callInfo)
+    {
+        return new ClientCookieMiddleware(this, _logger);
+    }
+
+    internal void UpdateCookies(IEnumerable<string> newCookieHeaderValues)
+    {
+        foreach (var headerValue in newCookieHeaderValues)
+        {
+            try
+            {
+                foreach (var parsedCookie in headerValue.ParseHeader())
+                {
+                    var nameLc = parsedCookie.Name.ToLower(CultureInfo.InvariantCulture);
+                    if (parsedCookie.IsExpired(headerValue))
+                    {
+                        Cookies.TryRemove(nameLc, out _);
+                    }
+                    else
+                    {
+                        Cookies[nameLc] = parsedCookie;
+                    }
+                }
+            }
+            catch (FormatException ex)
+            {
+                _logger.LogWarning(ex, "Skipping malformed Set-Cookie header: '{HeaderValue}'", headerValue);
+            }
+        }
+    }
+}

csharp/src/Apache.Arrow.Flight/Middleware/Extensions/CookieExtensions.cs

@@ -0,0 +1,104 @@
+// Licensed to the Apache Software Foundation (ASF) under one or more
+// contributor license agreements. See the NOTICE file distributed with
+// this work for additional information regarding copyright ownership.
+// The ASF licenses this file to You under the Apache License, Version 2.0
+// (the "License"); you may not use this file except in compliance with
+// the License.  You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+using System;
+using System.Collections.Generic;
+using System.Globalization;
+using System.Linq;
+using System.Net;
+
+namespace Apache.Arrow.Flight.Middleware.Extensions;
+
+public static class CookieExtensions
+{
+    public static IEnumerable<Cookie> ParseHeader(this string setCookieHeader)
+    {
+        if (string.IsNullOrWhiteSpace(setCookieHeader))
+            return System.Array.Empty<Cookie>();
+
+        var cookies = new List<Cookie>();
+
+        var segments = setCookieHeader.Split([';'], StringSplitOptions.RemoveEmptyEntries);
+        if (segments.Length == 0)
+            return cookies;
+
+        var nameValue = segments[0].Split(['='], 2);
+        if (nameValue.Length != 2 || string.IsNullOrWhiteSpace(nameValue[0]))
+            return cookies;
+
+        var name = nameValue[0].Trim();
+        var value = nameValue[1].Trim();
+        var cookie = new Cookie(name, value);
+
+        foreach (var segment in segments.Skip(1))
+        {
+            var kv = segment.Split(['='], 2, StringSplitOptions.RemoveEmptyEntries);
+            var key = kv[0].Trim().ToLowerInvariant();
+            var val = kv.Length > 1 ? kv[1] : null;
+
+            switch (key)
+            {
+                case "expires":
+                    if (!string.IsNullOrWhiteSpace(val))
+                    {
+                        if (DateTimeOffset.TryParseExact(val, "R", CultureInfo.InvariantCulture, DateTimeStyles.None, out var expiresRfc))
+                            cookie.Expires = expiresRfc.UtcDateTime;
+                        else if (DateTimeOffset.TryParse(val, out var expiresFallback))
+                            cookie.Expires = expiresFallback.UtcDateTime;
+                    }
+                    break;
+
+                case "max-age":
+                    if (int.TryParse(val, out var seconds))
+                        cookie.Expires = DateTime.UtcNow.AddSeconds(seconds);
+                    break;
+
+                case "domain":
+                    cookie.Domain = val;
+                    break;
+
+                case "path":
+                    cookie.Path = val;
+                    break;
+
+                case "secure":
+                    cookie.Secure = true;
+                    break;
+
+                case "httponly":
+                    cookie.HttpOnly = true;
+                    break;
+            }
+        }
+
+        cookies.Add(cookie);
+        return cookies;
+    }
+
+    public static bool IsExpired(this Cookie cookie, string rawHeader)
+    {
+        if (string.IsNullOrWhiteSpace(cookie?.Value))
+            return true;
+
+        // If raw header has Max-Age=0, consider it deleted
+        if (rawHeader?.IndexOf("Max-Age=0", StringComparison.OrdinalIgnoreCase) >= 0)
+            return true;
+
+        if (cookie.Expires != DateTime.MinValue && cookie.Expires <= DateTime.UtcNow)
+            return true;
+
+        return false;
+    }
+}